# SPDX-License-Identifier: AGPL-3.0-only
# Copyright (c) 2026 sol pbc

"""Media retention service for solstone journals.

Manages the lifecycle of raw media files (layer 1 captures) in journal segments.
Three retention modes:
- keep: retain raw media indefinitely and purge nothing (default)
- days: delete raw media after explicit N days, once processing is complete
- processed: delete raw media as soon as processing completes

Scope: raw media ONLY. Chronicle JSONL, derived outputs, talents/ directories,
and all other journal content persist indefinitely and are never touched by
retention. Days mode requires an explicit days value from config or CLI override;
without one, it purges nothing.

Safety invariant: never delete raw media from segments that haven't finished
processing. All completion checks must pass before any deletion.
"""

from __future__ import annotations

import hashlib
import json
import logging
import shutil
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Any

from solstone.apps.backup.copy import OFFLOAD_STALL_REASON_LABELS, OFFLOAD_STALLED_LEAD
from solstone.think.backup.state import merge_backup_config
from solstone.think.data_state import DataState, derive_modality_state
from solstone.think.media import AUDIO_EXTENSIONS as RAW_AUDIO_EXTENSIONS
from solstone.think.media import MEDIA_EXTENSIONS as RAW_MEDIA_EXTENSIONS
from solstone.think.media import VIDEO_EXTENSIONS as RAW_VIDEO_EXTENSIONS
from solstone.think.pruning_audit import AuditOutcome, write_prune_audit
from solstone.think.utils import day_dirs, get_journal, iter_segments

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Raw media file identification
# ---------------------------------------------------------------------------


def is_raw_media(path: Path) -> bool:
    """Check if a file is raw media (layer 1 capture).

    Raw media: every registry container format in ``solstone.think.media.FORMATS``
    — audio (*.flac, *.opus, *.ogg, *.m4a, *.mp3, *.wav), video (*.webm, *.mp4,
    *.mov), and still images (*.png, *.jpg, *.jpeg, *.heic, *.heif, *.gif,
    *.webp, *.tiff — e.g. mentra-live photos) — plus monitor_*_diff.png screen
    diffs (covered by the registry since images joined it; kept explicit).
    """
    if path.suffix.lower() in RAW_MEDIA_EXTENSIONS:
        return True
    if (
        path.suffix.lower() == ".png"
        and path.name.startswith("monitor_")
        and "_diff" in path.name
    ):
        return True
    return False


def get_raw_media_files(segment_path: Path) -> list[Path]:
    """Return all raw media files in a segment directory."""
    if not segment_path.is_dir():
        return []
    return [f for f in segment_path.iterdir() if f.is_file() and is_raw_media(f)]


# ---------------------------------------------------------------------------
# Completion detection (safety invariant)
# ---------------------------------------------------------------------------


@dataclass
class SegmentGate:
    """Deletion gate verdict for one segment."""

    verdict: str
    failed_files: dict[str, str] = field(default_factory=dict)
    completion_files: list[Path] = field(default_factory=list)


def _matching_extraction_files(
    segment_path: Path, patterns: tuple[str, ...]
) -> list[Path]:
    return sorted(
        {
            path
            for pattern in patterns
            for path in segment_path.glob(pattern)
            if path.is_file()
        }
    )


def _derive_extraction_file_state(
    seg_path: Path,
    jsonl_path: Path,
    *,
    modality: str,
    marker_key: str,
    has_raw: bool,
) -> str:
    """Read one extraction file strictly enough for irreversible deletion."""
    try:
        lines: list[str] = []
        with jsonl_path.open("r", encoding="utf-8") as handle:
            for line in handle:
                if not line.strip():
                    continue
                lines.append(line)
                if len(lines) == 2:
                    break
    except OSError:
        return "malformed"

    if not lines:
        return "malformed"

    try:
        header = json.loads(lines[0])
    except (json.JSONDecodeError, ValueError):
        return "malformed"
    if not isinstance(header, dict):
        return "malformed"

    record = header.get("_solstone_processing")
    if not isinstance(record, dict):
        record = None

    # Row 1 is always the metadata header. A stray marker key merged through
    # SEGMENT_META must never make a header-only file look chunk-bearing when
    # retention is deciding whether raw media can be irreversibly deleted.
    has_chunks = False
    if len(lines) == 2:
        try:
            first_chunk = json.loads(lines[1])
        except (json.JSONDecodeError, ValueError):
            first_chunk = None
        has_chunks = isinstance(first_chunk, dict) and marker_key in first_chunk

    return derive_modality_state(
        seg_path,
        modality,
        has_chunks=has_chunks,
        has_jsonl=True,
        has_raw=has_raw,
        record=record,
    )


def _collect_extraction_states(
    segment_path: Path,
    paths: list[Path],
    *,
    modality: str,
    marker_key: str,
    has_raw: bool,
    failed_files: dict[str, str],
    completion_files: list[Path],
) -> bool:
    incomplete = False
    for path in paths:
        state = _derive_extraction_file_state(
            segment_path,
            path,
            modality=modality,
            marker_key=marker_key,
            has_raw=has_raw,
        )
        if state == "malformed":
            failed_files[path.name] = "malformed"
        elif state in {DataState.FAILED.value, DataState.FAILED_FINAL.value}:
            failed_files[path.name] = state
        elif state in {DataState.PENDING.value, DataState.ANALYZING.value}:
            incomplete = True
        elif state in {DataState.ANALYZED.value, DataState.EMPTY.value}:
            completion_files.append(path)
        else:  # pragma: no cover - derive_modality_state owns this closed vocabulary.
            raise RuntimeError(f"unexpected {modality} data state for {path}: {state}")
    return incomplete


def resolve_segment_gate(segment_path: Path) -> SegmentGate:
    """Resolve whether a segment's raw media is safe to purge."""
    agents_dir = segment_path / "talents"
    incomplete = False
    failed_files: dict[str, str] = {}
    completion_files: list[Path] = []

    if agents_dir.is_dir():
        for f in agents_dir.iterdir():
            if f.is_file() and f.name.endswith("_active.jsonl"):
                incomplete = True
                break

    files = [f for f in segment_path.iterdir() if f.is_file()]
    file_suffixes = {f.suffix.lower() for f in files}
    has_audio_raw = bool(file_suffixes & RAW_AUDIO_EXTENSIONS)
    has_video_raw = bool(file_suffixes & RAW_VIDEO_EXTENSIONS)
    audio_extracts = _matching_extraction_files(
        segment_path, ("audio.jsonl", "*_audio.jsonl")
    )
    screen_extracts = _matching_extraction_files(
        segment_path, ("screen.jsonl", "*_screen.jsonl")
    )

    if has_audio_raw and not audio_extracts:
        incomplete = True

    # monitor_*_diff.png files are raw media but have no extraction record. They
    # ride the whole-segment gate and are only deleted when audio/video checks pass.
    if has_video_raw and not screen_extracts:
        incomplete = True

    speaker_labels = agents_dir / "speaker_labels.json"
    if ".npz" in file_suffixes:
        if not agents_dir.is_dir() or not speaker_labels.exists():
            incomplete = True

    incomplete = (
        _collect_extraction_states(
            segment_path,
            audio_extracts,
            modality="audio",
            marker_key="start",
            has_raw=has_audio_raw,
            failed_files=failed_files,
            completion_files=completion_files,
        )
        or incomplete
    )
    incomplete = (
        _collect_extraction_states(
            segment_path,
            screen_extracts,
            modality="screen",
            marker_key="timestamp",
            has_raw=has_video_raw,
            failed_files=failed_files,
            completion_files=completion_files,
        )
        or incomplete
    )

    if failed_files:
        return SegmentGate("failed", failed_files=failed_files)
    if speaker_labels.exists():
        completion_files.append(speaker_labels)
    if incomplete:
        return SegmentGate("incomplete", completion_files=completion_files)
    return SegmentGate("eligible", completion_files=completion_files)


# ---------------------------------------------------------------------------
# Retention configuration
# ---------------------------------------------------------------------------


@dataclass
class RetentionPolicy:
    """Retention policy for a single scope (global or per-stream)."""

    mode: str = "keep"  # "keep", "days", or "processed"
    days: int | None = None

    def is_eligible(self, segment_age_days: int) -> bool:
        """Check if a segment's raw media should be purged under this policy."""
        if self.mode == "keep":
            return False
        if self.mode == "processed":
            return True
        if self.mode == "days" and self.days is not None:
            return segment_age_days >= self.days
        return False


@dataclass
class RetentionConfig:
    """Retention configuration from journal.json."""

    default: RetentionPolicy = field(default_factory=RetentionPolicy)
    per_stream: dict[str, RetentionPolicy] = field(default_factory=dict)

    def policy_for_stream(self, stream: str) -> RetentionPolicy:
        """Return the effective policy for a stream."""
        return self.per_stream.get(stream, self.default)


def load_retention_config() -> RetentionConfig:
    """Load retention configuration from journal.json."""
    from solstone.think.utils import get_config

    config = get_config()
    retention = config.get("retention", {})

    mode = retention.get("raw_media", "keep")
    days = retention.get("raw_media_days", None)
    default = RetentionPolicy(mode=mode, days=days)

    per_stream: dict[str, RetentionPolicy] = {}
    for stream_name, stream_config in retention.get("per_stream", {}).items():
        per_stream[stream_name] = RetentionPolicy(
            mode=stream_config.get("raw_media", mode),
            days=stream_config.get("raw_media_days", days),
        )

    return RetentionConfig(default=default, per_stream=per_stream)


# ---------------------------------------------------------------------------
# Storage summary
# ---------------------------------------------------------------------------


def _human_bytes(size: int) -> str:
    """Format byte count as human-readable string."""
    n = float(size)
    for unit in ("B", "KB", "MB", "GB", "TB"):
        if abs(n) < 1024:
            if unit == "B":
                return f"{int(n)} B"
            return f"{n:.1f} {unit}"
        n /= 1024
    return f"{n:.1f} PB"


@dataclass
class StorageSummary:
    """Storage usage summary for a journal."""

    raw_media_bytes: int = 0
    derived_bytes: int = 0
    total_segments: int = 0
    segments_with_raw: int = 0
    segments_purged: int = 0

    @property
    def raw_media_human(self) -> str:
        return _human_bytes(self.raw_media_bytes)

    @property
    def derived_human(self) -> str:
        return _human_bytes(self.derived_bytes)


def compute_storage_summary() -> StorageSummary:
    """Compute storage summary across all journal segments."""
    summary = StorageSummary()

    for day_name in sorted(day_dirs().keys()):
        for _stream, _seg_key, seg_path in iter_segments(day_name):
            summary.total_segments += 1

            raw_files = get_raw_media_files(seg_path)
            summary.raw_media_bytes += sum(f.stat().st_size for f in raw_files)

            if raw_files:
                summary.segments_with_raw += 1
            elif (seg_path / "audio.jsonl").exists() or (
                seg_path / "screen.jsonl"
            ).exists():
                summary.segments_purged += 1

            for f in seg_path.rglob("*"):
                if f.is_file() and not is_raw_media(f):
                    summary.derived_bytes += f.stat().st_size

    return summary


def check_storage_health(
    summary: StorageSummary,
    journal_path: str | Path,
    config: dict | None = None,
) -> list[dict]:
    """Check storage health against configured thresholds.

    Parameters
    ----------
    summary
        Pre-computed storage summary (avoids recomputation).
    journal_path
        Journal root path, used for disk usage check.
    config
        Full journal config dict. Loaded via get_config() if not provided.

    Returns
    -------
    list[dict]
        List of warning dicts. Empty when healthy.
    """
    if config is None:
        from solstone.think.utils import get_config

        config = get_config()

    retention = config.get("retention", {})
    keep_mode_nudge = (
        " you're currently set to always retain original media — consider "
        "choosing a retention value in settings to free up disk space."
    )
    always_retain_enabled = retention.get("raw_media", "keep") == "keep"
    warnings = []

    # Check disk usage percentage
    disk_threshold = retention.get("storage_warning_disk_percent", 80)
    if disk_threshold is not None:
        try:
            usage = shutil.disk_usage(str(journal_path))
            disk_percent = round(usage.used / usage.total * 100, 1)
            if disk_percent >= disk_threshold:
                message = (
                    f"Disk is {disk_percent}% full (threshold: {disk_threshold}%). "
                    "Consider adjusting retention settings or running Clean Up Now "
                    "to free space."
                )
                if always_retain_enabled:
                    message += keep_mode_nudge
                warnings.append(
                    {
                        "level": "warning",
                        "type": "disk_percent",
                        "message": message,
                        "current": disk_percent,
                        "threshold": disk_threshold,
                    }
                )
        except OSError:
            pass

    # Check raw media size
    raw_media_gb_threshold = retention.get("storage_warning_raw_media_gb")
    if raw_media_gb_threshold is not None:
        raw_media_gb = round(summary.raw_media_bytes / (1024**3), 2)
        if raw_media_gb >= raw_media_gb_threshold:
            message = (
                f"Raw media is {raw_media_gb} GB (threshold: {raw_media_gb_threshold} GB). "
                "Consider adjusting retention settings or running Clean Up Now "
                "to free space."
            )
            if always_retain_enabled:
                message += keep_mode_nudge
            warnings.append(
                {
                    "level": "warning",
                    "type": "raw_media_gb",
                    "message": message,
                    "current": raw_media_gb,
                    "threshold": raw_media_gb_threshold,
                }
            )

    backup = merge_backup_config(config)
    offload = backup["offload"]
    last_offload = backup["last_offload"]
    if offload.get("enabled") is True and last_offload.get("status") == "stalled":
        reason = last_offload.get("reason")
        reason_label = (
            OFFLOAD_STALL_REASON_LABELS.get(reason) if isinstance(reason, str) else None
        )
        warnings.append(
            {
                "level": "warning",
                "type": "offload_stalled",
                "message": " ".join(
                    part for part in (OFFLOAD_STALLED_LEAD, reason_label) if part
                ),
                "current": None,
                "threshold": None,
            }
        )

    return warnings


# ---------------------------------------------------------------------------
# Retention purge
# ---------------------------------------------------------------------------


@dataclass
class PurgeResult:
    """Result of a purge operation."""

    files_deleted: int = 0
    bytes_freed: int = 0
    segments_processed: int = 0
    segments_skipped_incomplete: int = 0
    segments_skipped_policy: int = 0
    segments_blocked_failed: int = 0
    blocked_failed_details: list[dict[str, Any]] = field(default_factory=list)
    partial_error: bool = False
    details: list[dict[str, Any]] = field(default_factory=list)


def purge(
    *,
    older_than_days: int | None = None,
    stream_filter: str | None = None,
    dry_run: bool = False,
    config: RetentionConfig | None = None,
) -> PurgeResult:
    """Run retention purge across the journal.

    Parameters
    ----------
    older_than_days
        Override: purge raw media older than this many days.
        If None, uses the configured retention policy.
    stream_filter
        Only process segments from this stream.
    dry_run
        If True, report what would be deleted without deleting.
    config
        Retention config. If None, loads from journal.json.
    """
    if config is None:
        config = load_retention_config()

    result = PurgeResult()
    today = datetime.now().date()
    journal_path = Path(get_journal())

    for day_name in sorted(day_dirs().keys()):
        try:
            day_date = datetime.strptime(day_name, "%Y%m%d").date()
        except ValueError:
            continue
        age_days = (today - day_date).days

        for stream_name, seg_key, seg_path in iter_segments(day_name):
            if stream_filter and stream_name != stream_filter:
                continue

            raw_files = get_raw_media_files(seg_path)
            if not raw_files:
                continue

            result.segments_processed += 1

            gate = resolve_segment_gate(seg_path)
            if gate.verdict == "failed":
                result.segments_blocked_failed += 1
                blocked_detail = {
                    "day": day_name,
                    "stream": stream_name,
                    "segment": seg_key,
                    "files": gate.failed_files,
                }
                result.blocked_failed_details.append(blocked_detail)
                logger.warning(
                    "retention: blocked purge - extraction failed: %s/%s/%s files=%s",
                    day_name,
                    stream_name,
                    seg_key,
                    gate.failed_files,
                )
                continue

            # Safety invariant: never delete from incomplete segments
            if gate.verdict == "incomplete":
                result.segments_skipped_incomplete += 1
                logger.debug(
                    "Skipping incomplete: %s/%s/%s", day_name, stream_name, seg_key
                )
                continue
            if gate.verdict != "eligible":
                raise RuntimeError(f"unexpected retention gate verdict: {gate.verdict}")

            # Check eligibility
            if older_than_days is not None:
                eligible = age_days >= older_than_days
            else:
                policy = config.policy_for_stream(stream_name)
                eligible = policy.is_eligible(age_days)

            if not eligible:
                result.segments_skipped_policy += 1
                continue

            # Delete raw media
            segment_bytes = 0
            segment_files = []
            for f in raw_files:
                size = f.stat().st_size
                digest = hashlib.sha256()
                with open(f, "rb") as handle:
                    while chunk := handle.read(64 * 1024):
                        digest.update(chunk)
                hex_digest = digest.hexdigest()
                segment_bytes += size
                segment_files.append(
                    {"name": f.name, "bytes": size, "hash": hex_digest}
                )
                if not dry_run:
                    f.unlink()
                    logger.info("Deleted: %s (%s)", f, _human_bytes(size))

            processed_at = None
            if gate.completion_files:
                latest_mtime = max(f.stat().st_mtime for f in gate.completion_files)
                processed_at = datetime.fromtimestamp(latest_mtime).isoformat()

            result.files_deleted += len(raw_files)
            result.bytes_freed += segment_bytes
            result.details.append(
                {
                    "day": day_name,
                    "stream": stream_name,
                    "segment": seg_key,
                    "files": segment_files,
                    "bytes_freed": segment_bytes,
                    "processed_at": processed_at,
                }
            )

            if not dry_run and segment_files:
                outcome = _write_segment_prune_audit(
                    journal_path,
                    day_name,
                    stream_name,
                    seg_key,
                    segment_files,
                    segment_bytes,
                    processed_at,
                    older_than_days,
                )
                for audit_day, error in outcome.per_day_failures.items():
                    result.partial_error = True
                    logger.warning(
                        "retention: failed to append pruning task log for %s: %s",
                        audit_day,
                        error,
                    )
                if outcome.global_record_error is not None:
                    result.partial_error = True
                    logger.warning(
                        "retention: failed to append pruning run record: %s",
                        outcome.global_record_error,
                    )

    if not dry_run and (result.files_deleted > 0 or result.segments_blocked_failed > 0):
        _write_retention_log(journal_path, result)

    return result


def _write_segment_prune_audit(
    journal_path: Path,
    day: str,
    stream: str,
    segment: str,
    files: list[dict[str, Any]],
    bytes_freed: int,
    processed_at: str | None,
    older_than_days: int | None,
) -> AuditOutcome:
    run_record = {
        "timestamp": datetime.now().isoformat(),
        "kind": "raw_media",
        "dry_run": False,
        "days": older_than_days,
        "day": day,
        "stream": stream,
        "segment": segment,
        "files": files,
        "bytes_freed": bytes_freed,
        "processed_at": processed_at,
    }
    message = (
        f"raw-media retention: pruned {len(files)} raw media file(s) "
        f"({_human_bytes(bytes_freed)}) from segment {stream}/{segment}"
    )
    return write_prune_audit(
        journal_path,
        kind="raw_media",
        run_record=run_record,
        per_day_messages={day: message},
    )


def _write_retention_log(journal_path: Path, result: PurgeResult) -> None:
    """Append retention activity to health/retention.log."""
    health_dir = journal_path / "health"
    health_dir.mkdir(parents=True, exist_ok=True)
    log_path = health_dir / "retention.log"

    entry = {
        "timestamp": datetime.now().isoformat(),
        "files_deleted": result.files_deleted,
        "bytes_freed": result.bytes_freed,
        "segments_processed": result.segments_processed,
        "segments_skipped_incomplete": result.segments_skipped_incomplete,
        "segments_blocked_failed": result.segments_blocked_failed,
        "blocked_failed_details": result.blocked_failed_details,
        "partial_error": result.partial_error,
        "details": result.details,
    }

    with open(log_path, "a", encoding="utf-8") as f:
        f.write(json.dumps(entry) + "\n")
