"""
Silent trajectory logging for the evolution pipeline.

Persists every completed agent run (messages, tool calls, results,
final reply) to encrypted JSONL under ~/.mugi/evolution/trajectories/.
Each line is Fernet-encrypted using the same Keychain key as chat_storage,
so trajectory data has the same at-rest protection as chat messages.

Called from agent_loop.run_raw_agent; failures are silent.
"""
from __future__ import annotations

import base64
import json
import os
from datetime import datetime, timedelta

_EVOLUTION_DIR = os.path.join(os.path.expanduser("~"), ".mugi", "evolution")
TRAJECTORIES_DIR = os.path.join(_EVOLUTION_DIR, "trajectories")

_DIR_MODE = 0o700
_FILE_MODE = 0o600
_MAX_MESSAGES_CHARS = 50_000
_MAX_REPLY_CHARS = 50_000
_MIN_MESSAGES_FOR_LOG = 4

TRAJECTORY_RETENTION_DAYS = 90


def _log_evolution(msg: str, level: str = "WARNING") -> None:
    try:
        from tools._logging import _debug_log
        _debug_log(f"[evolution/trajectory_logger] {msg}", level)
    except Exception:
        pass


def _get_fernet():
    """Reuse the Fernet instance from chat_storage (same Keychain key)."""
    from chat_storage import _get_fernet
    return _get_fernet()


def _encrypt_line(plaintext: str) -> str:
    """Encrypt a string and return a base64 token suitable for one JSONL line."""
    fernet = _get_fernet()
    return fernet.encrypt(plaintext.encode("utf-8")).decode("ascii")


def decrypt_line(encrypted_line: str) -> str:
    """Decrypt one JSONL line back to the original JSON string."""
    fernet = _get_fernet()
    return fernet.decrypt(encrypted_line.encode("ascii")).decode("utf-8")


def _ensure_trajectories_dir() -> str:
    os.makedirs(TRAJECTORIES_DIR, mode=_DIR_MODE, exist_ok=True)
    try:
        os.chmod(TRAJECTORIES_DIR, _DIR_MODE)
    except OSError:
        pass
    return TRAJECTORIES_DIR


def _truncate(text: str, max_chars: int) -> str:
    if len(text) <= max_chars:
        return text
    return text[:max_chars] + "\n[truncated]"


def _sanitize_messages(messages: list) -> list:
    """Produce a JSON-safe copy of the messages list, truncating large content."""
    sanitized = []
    total_chars = 0
    for m in messages:
        if not isinstance(m, dict):
            continue
        entry = {}
        for k, v in m.items():
            entry[k] = v
        sanitized.append(entry)
        total_chars += len(json.dumps(entry, default=str))
        if total_chars > _MAX_MESSAGES_CHARS:
            break
    return sanitized


def _extract_tool_calls(messages: list) -> list[dict]:
    calls: list[dict] = []
    for m in messages:
        if not isinstance(m, dict) or m.get("role") != "assistant":
            continue
        for tc in (m.get("tool_calls") or []):
            if not isinstance(tc, dict):
                continue
            fn = tc.get("function") or {}
            if not isinstance(fn, dict):
                continue
            name = str(fn.get("name") or "").strip()
            if not name:
                continue
            calls.append(
                {
                    "name": name,
                    "arguments": _truncate(str(fn.get("arguments") or ""), 500),
                }
            )
    return calls


def _purge_old_trajectories() -> None:
    """Delete trajectory files older than TRAJECTORY_RETENTION_DAYS."""
    try:
        if not os.path.isdir(TRAJECTORIES_DIR):
            return
        cutoff = datetime.now() - timedelta(days=TRAJECTORY_RETENTION_DAYS)
        for fname in os.listdir(TRAJECTORIES_DIR):
            if not fname.endswith(".jsonl"):
                continue
            # Extract date from filename: trajectories_YYYY-MM-DD.jsonl
            try:
                date_part = fname.replace("trajectories_", "").replace(".jsonl", "")
                file_date = datetime.strptime(date_part, "%Y-%m-%d")
                if file_date < cutoff:
                    os.unlink(os.path.join(TRAJECTORIES_DIR, fname))
            except (ValueError, OSError):
                continue
    except Exception as e:
        _log_evolution(f"purge_old_trajectories failed: {e}")


def persist_trajectory(
    thread_id: str,
    messages: list,
    response_buffer: str,
    model_id: str | None = None,
    active_drives_snapshot: list[dict] | None = None,
) -> None:
    """Append one encrypted trajectory record to today's JSONL file.

    Each line is Fernet-encrypted so trajectory data is protected at rest.
    Safe to call from any context; all exceptions are swallowed so
    logging never blocks or crashes the agent.
    """
    try:
        if len(messages) < _MIN_MESSAGES_FOR_LOG:
            return

        _ensure_trajectories_dir()

        # Opportunistically purge old files (cheap stat-only check)
        _purge_old_trajectories()

        record = {
            "thread_id": thread_id or "unknown",
            "timestamp_iso": datetime.now().isoformat(),
            "model": model_id or "",
            "messages": _sanitize_messages(messages),
            "tool_calls": _extract_tool_calls(messages),
            "final_reply": _truncate(response_buffer or "", _MAX_REPLY_CHARS),
        }
        if active_drives_snapshot:
            record["active_drives_at_start"] = active_drives_snapshot

        plaintext = json.dumps(record, default=str)
        encrypted = _encrypt_line(plaintext)

        today = datetime.now().strftime("%Y-%m-%d")
        path = os.path.join(TRAJECTORIES_DIR, f"trajectories_{today}.jsonl")

        fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_APPEND, _FILE_MODE)
        try:
            with os.fdopen(fd, "a") as f:
                f.write(encrypted + "\n")
        except Exception:
            try:
                os.close(fd)
            except Exception:
                pass
    except Exception as e:
        _log_evolution(f"persist_trajectory failed: {e}", "ERROR")


def load_trajectory_records(since_iso: str | None = None) -> list[dict]:
    """Load and decrypt trajectory records, optionally filtered by timestamp.

    Used by grpo_runner to build training datasets.
    """
    records = []
    if not os.path.isdir(TRAJECTORIES_DIR):
        return records
    since_dt = None
    if since_iso:
        try:
            since_dt = datetime.fromisoformat(since_iso)
        except Exception:
            pass
    for fname in sorted(os.listdir(TRAJECTORIES_DIR)):
        if not fname.endswith(".jsonl"):
            continue
        fpath = os.path.join(TRAJECTORIES_DIR, fname)
        try:
            with open(fpath, "r", encoding="utf-8") as f:
                for line in f:
                    line = line.strip()
                    if not line:
                        continue
                    try:
                        plaintext = decrypt_line(line)
                        record = json.loads(plaintext)
                        if since_dt:
                            ts = record.get("timestamp_iso", "")
                            try:
                                if datetime.fromisoformat(ts) <= since_dt:
                                    continue
                            except Exception:
                                pass
                        records.append(record)
                    except Exception:
                        continue
        except Exception:
            continue
    return records


def count_trajectory_records(since_iso: str | None = None) -> int:
    """Count trajectory records, optionally filtered by timestamp.

    Used by the scheduler to check data thresholds without loading all records.
    """
    count = 0
    if not os.path.isdir(TRAJECTORIES_DIR):
        return 0
    since_dt = None
    if since_iso:
        try:
            since_dt = datetime.fromisoformat(since_iso)
        except Exception:
            pass
    for fname in os.listdir(TRAJECTORIES_DIR):
        if not fname.endswith(".jsonl"):
            continue
        fpath = os.path.join(TRAJECTORIES_DIR, fname)
        try:
            with open(fpath, "r", encoding="utf-8") as f:
                for line in f:
                    line = line.strip()
                    if not line:
                        continue
                    if since_dt:
                        try:
                            plaintext = decrypt_line(line)
                            record = json.loads(plaintext)
                            ts = record.get("timestamp_iso", "")
                            if datetime.fromisoformat(ts) <= since_dt:
                                continue
                        except Exception:
                            pass
                    count += 1
        except Exception:
            continue
    return count
