# tools/conversation_trace.py — Conversation trace logging for debugging.
#
# When enabled via [debug] trace_conversations in mugi_config.toml, logs each
# agent turn (LLM request/response, tool calls/results) to JSONL files under
# ~/.mugi/logs/traces/. Use scripts/replay_trace.py to replay a trace without
# calling the LLM.
from __future__ import annotations

import json
import os
import hashlib
import threading
import time
from typing import Any

_TRACE_LOCK = threading.Lock()
_TRACE_ENABLED: bool | None = None
_TRACE_DIR: str | None = None


def reset_trace_config() -> None:
    """Clear cached trace config so the next read reflects current config."""
    global _TRACE_ENABLED, _TRACE_DIR
    _TRACE_ENABLED = None
    _TRACE_DIR = None


def _get_trace_config() -> tuple[bool, str]:
    """Return (enabled, trace_dir). Reads config once per process."""
    global _TRACE_ENABLED, _TRACE_DIR
    if _TRACE_ENABLED is not None:
        return _TRACE_ENABLED, _TRACE_DIR or ""
    try:
        from core.config import get_config
        cfg = get_config()
        debug = cfg.get("debug") or {}
        _TRACE_ENABLED = bool(debug.get("enabled", False)) and bool(debug.get("trace_conversations", False))
        raw_dir = (debug.get("trace_dir") or "").strip()
        if raw_dir:
            _TRACE_DIR = os.path.abspath(os.path.expanduser(raw_dir))
        else:
            from persistence import MUGI_LOGS_DIR
            _TRACE_DIR = os.path.join(MUGI_LOGS_DIR, "traces")
    except Exception:
        _TRACE_ENABLED = False
        _TRACE_DIR = ""
    return _TRACE_ENABLED, _TRACE_DIR or ""


def _ensure_trace_dir() -> str | None:
    enabled, trace_dir = _get_trace_config()
    if not enabled or not trace_dir:
        return None
    try:
        os.makedirs(trace_dir, exist_ok=True)
        return trace_dir
    except OSError:
        return None


def _trace_file_path(thread_id: str) -> str | None:
    base = _ensure_trace_dir()
    if not base:
        return None
    safe = "".join(c if c.isalnum() or c in "-_" else "_" for c in thread_id)
    return os.path.join(base, f"{safe}.jsonl")


def _write_trace_entry(thread_id: str, event: str, data: dict) -> None:
    path = _trace_file_path(thread_id)
    if not path:
        return
    payload = {
        "event": event,
        "thread_id": thread_id,
        "timestamp": time.time(),
        **data,
    }
    with _TRACE_LOCK:
        try:
            with open(path, "a", encoding="utf-8") as f:
                f.write(json.dumps(payload, ensure_ascii=False) + "\n")
        except OSError:
            pass


def _truncate_for_trace(obj: Any, max_chars: int = 2000) -> Any:
    """Truncate strings in nested structures for trace (avoid huge payloads)."""
    if isinstance(obj, str):
        if len(obj) <= max_chars:
            return obj
        return obj[:max_chars] + f"... [truncated, {len(obj)} chars total]"
    if isinstance(obj, dict):
        return {k: _truncate_for_trace(v, max_chars) for k, v in obj.items()}
    if isinstance(obj, list):
        return [_truncate_for_trace(v, max_chars) for v in obj]
    return obj


def _string_content_len(value: Any) -> int:
    if isinstance(value, str):
        return len(value)
    if isinstance(value, list):
        total = 0
        for part in value:
            if isinstance(part, dict):
                for key in ("text", "content"):
                    v = part.get(key)
                    if isinstance(v, str):
                        total += len(v)
            elif isinstance(part, str):
                total += len(part)
        return total
    if value is None:
        return 0
    return len(str(value))


def _json_len(value: Any) -> int:
    try:
        return len(json.dumps(value, ensure_ascii=False, default=str))
    except Exception:
        return len(str(value))


def is_trace_enabled() -> bool:
    enabled, _ = _get_trace_config()
    return enabled


def emit_turn_start(thread_id: str, round_: int, message_count: int, model: str) -> None:
    """Log start of a new agent turn."""
    _write_trace_entry(thread_id, "turn_start", {
        "round": round_,
        "message_count": message_count,
        "model": model,
    })


def emit_llm_request(
    thread_id: str,
    round_: int,
    messages: list,
    create_kwargs: dict,
    *,
    budget_messages_est: int | None = None,
    budget_tools_est: int | None = None,
    budget_total_est: int | None = None,
    pressure_ratio: float | None = None,
    cap_tokens: int | None = None,
    trim_mode: str | None = None,
) -> None:
    """Log LLM request (messages truncated for size).

    Optional budget_* fields come from utils.context_token_budget when trace is enabled,
    so operators can correlate failures with underestimated tool schema payload size.
    """
    def _preview_msg(m: dict) -> dict:
        role = m.get("role")
        content = m.get("content")
        if role == "user" and isinstance(content, str) and content.lstrip().startswith("<tool_response>"):
            return {"role": role, "content_preview": "[tool_response omitted]"}
        return {"role": role, "content_preview": (str(content)[:500] if content else None)}

    api_msgs = _truncate_for_trace(
        [_preview_msg(m) for m in (messages or [])]
    )
    request_messages = messages or []
    content_chars_by_role: dict[str, int] = {}
    total_content_chars = 0
    for m in request_messages:
        role = str(m.get("role") or "")
        content_len = _string_content_len(m.get("content"))
        total_content_chars += content_len
        content_chars_by_role[role] = content_chars_by_role.get(role, 0) + content_len
    tool_schemas = create_kwargs.get("tools") or []
    tool_schema_count = len(tool_schemas) if isinstance(tool_schemas, list) else 0
    tool_schema_chars = _json_len(tool_schemas) if tool_schema_count else 0
    payload: dict[str, Any] = {
        "round": round_,
        "messages_preview": api_msgs,
        "model": create_kwargs.get("model"),
        "max_tokens": create_kwargs.get("max_tokens"),
        "temperature": create_kwargs.get("temperature"),
        "message_count": len(request_messages),
        "content_chars_total": total_content_chars,
        "content_chars_by_role": content_chars_by_role,
        "system_prompt_chars": sum(
            _string_content_len(m.get("content"))
            for m in request_messages
            if m.get("role") == "system"
        ),
        "tool_schema_count": tool_schema_count,
        "tool_schema_chars": tool_schema_chars,
        "request_chars_estimate": total_content_chars + tool_schema_chars,
        "request_tokens_estimate": int((total_content_chars + tool_schema_chars) / 4),
    }
    if budget_messages_est is not None:
        payload["budget_messages_est"] = budget_messages_est
    if budget_tools_est is not None:
        payload["budget_tools_est"] = budget_tools_est
    if budget_total_est is not None:
        payload["budget_total_est"] = budget_total_est
    if pressure_ratio is not None:
        payload["pressure_ratio"] = pressure_ratio
    if cap_tokens is not None:
        payload["cap_tokens"] = cap_tokens
    if trim_mode:
        payload["trim_mode"] = trim_mode
    _write_trace_entry(thread_id, "llm_request", payload)


def emit_llm_response(
    thread_id: str,
    round_: int,
    assistant_content: str,
    tool_calls: list,
    metadata: dict | None = None,
) -> None:
    """Log LLM response (content + tool calls)."""
    payload = {
        "round": round_,
        "assistant_content_preview": _truncate_for_trace(assistant_content or "", 1500),
        "assistant_content_chars": len(assistant_content or ""),
        "tool_call_count": len(tool_calls or []),
        "tool_calls": _truncate_for_trace([{"name": tc.get("name"), "id": tc.get("id"), "arguments_preview": (tc.get("arguments") or "")[:500]} for tc in (tool_calls or [])]),
    }
    if isinstance(metadata, dict):
        payload.update(_truncate_for_trace(metadata, 500))
    _write_trace_entry(thread_id, "llm_response", payload)


def emit_tool_call(thread_id: str, round_: int, tool_name: str, args: str, result: str, duration_sec: float | None = None) -> None:
    """Log a tool invocation and its result with bounded argument previews."""
    _args = args or ""
    _args_preview = _truncate_for_trace(_args, 500)
    _args_for_replay = _args if len(_args) <= 2000 else (_args[:2000] + f"...[truncated sha12={_hash_for_trace(_args)}]")
    _args_hashed = _hash_for_trace(_args)
    _write_trace_entry(thread_id, "tool_call", {
        "round": round_,
        "tool_name": tool_name,
        "args": _args_for_replay,
        "args_preview": _args_preview,
        "args_sha12": _args_hashed,
        "args_truncated": len(_args) > 500,
        "result_preview": _truncate_for_trace(result or "", 1500),
        "duration_sec": round(duration_sec, 3) if duration_sec is not None else None,
    })


def emit_validation_guard(
    thread_id: str,
    round_: int,
    tool_name: str,
    missing_args: list[str],
    provided_args: list[str],
    normalized_args: list[str],
    retry_count: int,
) -> None:
    """Log preflight validation guard events (missing required tool args)."""
    _write_trace_entry(thread_id, "validation_guard", {
        "round": round_,
        "tool_name": tool_name,
        "missing_args": _truncate_for_trace(missing_args, 200),
        "provided_args": _truncate_for_trace(provided_args, 200),
        "normalized_args": _truncate_for_trace(normalized_args, 200),
        "retry_count": retry_count,
    })


def emit_turn_end(thread_id: str, round_: int, final_content: str | None, reason: str) -> None:
    """Log end of turn (final reply or loop break reason)."""
    _write_trace_entry(thread_id, "turn_end", {
        "round": round_,
        "final_content_preview": _truncate_for_trace(final_content or "", 1000),
        "reason": reason,
    })


def _hash_for_trace(value: str) -> str:
    """Return a short stable hash for sensitive trace values."""
    if not isinstance(value, str):
        return ""
    return hashlib.sha256(value.encode("utf-8")).hexdigest()[:12]


def emit_guardrail_event(
    thread_id: str,
    round_: int,
    guard_type: str,
    action: str,
    *,
    details: dict[str, Any] | None = None,
) -> None:
    """Log structured guard/recovery events with privacy-safe defaults."""
    safe_details = _truncate_for_trace(details or {}, 500)
    payload: dict[str, Any] = {
        "round": round_,
        "guard_type": (guard_type or "").strip(),
        "action": (action or "").strip(),
        "details": safe_details,
    }
    try:
        from utils.agent_loop_pressure_context import get_current_pressure_ratio

        _pr = get_current_pressure_ratio()
        if _pr is not None:
            _pf = float(_pr)
            payload["context_pressure_ratio"] = _pf
            # Match emit_llm_request field name for operator queries / exports.
            payload["pressure_ratio"] = _pf
    except Exception:
        pass
    _write_trace_entry(thread_id, "guardrail_event", payload)


def emit_sanitize_event(
    thread_id: str,
    round_: int,
    stage: str,
    *,
    removed_chars: int,
    removed_sections: dict[str, int] | None = None,
) -> None:
    """Log sanitize delta counters without storing raw content."""
    _write_trace_entry(thread_id, "sanitize_event", {
        "round": round_,
        "stage": (stage or "").strip(),
        "removed_chars": max(0, int(removed_chars or 0)),
        "removed_sections": _truncate_for_trace(removed_sections or {}, 200),
    })


def emit_drive_event(
    thread_id: str,
    event_type: str,
    *,
    drive_id: str = "",
    details: dict[str, Any] | None = None,
) -> None:
    """Log structured drive lifecycle/initiative events."""
    _write_trace_entry(thread_id, "drive_event", {
        "event_type": (event_type or "").strip(),
        "drive_id": (drive_id or "").strip(),
        "details": _truncate_for_trace(details or {}, 500),
    })


def emit_run_start(thread_id: str, subagent_mode: bool, idle_tick: bool = False) -> None:
    """Log start of run_raw_agent."""
    _write_trace_entry(thread_id, "run_start", {"subagent_mode": subagent_mode, "idle_tick": idle_tick})


def emit_run_end(thread_id: str) -> None:
    """Log end of run_raw_agent."""
    _write_trace_entry(thread_id, "run_end", {})


def emit_subagent_routing(
    thread_id: str,
    tool_name: str,
    profile: str,
    elapsed_seconds: float,
    success: bool,
    timed_out: bool,
) -> None:
    """Log a subagent routing decision and its outcome."""
    _write_trace_entry(thread_id, "subagent_routing", {
        "tool_name": tool_name,
        "profile": profile,
        "elapsed_seconds": round(elapsed_seconds, 3),
        "success": success,
        "timed_out": timed_out,
    })
