"""Codex session adapter for normalized viewer and index records."""

from __future__ import annotations

from datetime import datetime
from pathlib import Path

import json

from lerim.adapters.base import SessionRecord, ViewerMessage, ViewerSession
from lerim.adapters.common import (
    count_non_empty_files,
    in_window,
    load_jsonl_dict_lines,
    parse_timestamp,
)


def compact_trace(raw_text: str) -> str:
    """Remove non-conversational noise and strip tool outputs from Codex trace JSONL.

    Drops: turn_context lines (full codebase snapshots).
    Strips: base_instructions from session_meta.
    Clears: function_call_output content (replaced with size descriptor).
    Clears: reasoning / agent_reasoning content (replaced with size descriptor).
    """
    kept: list[str] = []
    for line in raw_text.split("\n"):
        stripped = line.strip()
        if not stripped:
            continue
        try:
            obj = json.loads(stripped)
        except json.JSONDecodeError:
            kept.append(line)
            continue
        if obj.get("type") == "turn_context":
            continue
        payload = obj.get("payload")
        if not isinstance(payload, dict):
            kept.append(line)
            continue
        line_type = obj.get("type")
        if line_type == "session_meta":
            payload.pop("base_instructions", None)
        elif line_type == "response_item":
            ptype = payload.get("type")
            if ptype == "function_call_output":
                output = payload.get("output", "")
                payload["output"] = f"[cleared: {len(output)} chars]"
            elif ptype == "reasoning":
                content = payload.get("content", [])
                total = (
                    sum(len(c.get("text", "")) for c in content if isinstance(c, dict))
                    if isinstance(content, list)
                    else len(str(content))
                )
                payload["content"] = f"[reasoning cleared: {total} chars]"
        elif line_type == "event_msg":
            if payload.get("type") == "agent_reasoning":
                msg = payload.get("message", "")
                payload["message"] = f"[reasoning cleared: {len(msg)} chars]"
        kept.append(json.dumps(obj, ensure_ascii=False))
    return "\n".join(kept) + "\n"


def _default_cache_dir() -> Path:
    """Return the default cache directory for compacted Codex JSONL files."""
    return Path("~/.lerim/cache/codex").expanduser()


def default_path() -> Path | None:
    """Return the default Codex session trace directory."""
    return Path("~/.codex/sessions/").expanduser()


def count_sessions(path: Path) -> int:
    """Count readable non-empty Codex session JSONL files."""
    return count_non_empty_files(path, "*.jsonl")


def _extract_message_text(content: object) -> str | None:
    """Normalize message payload content to plain text."""
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        parts: list[str] = []
        for block in content:
            if isinstance(block, dict):
                text = block.get("text")
                if isinstance(text, str):
                    parts.append(text)
        if parts:
            return "\n".join(parts)
    return None


def find_session_path(session_id: str, traces_dir: Path | None = None) -> Path | None:
    """Find a Codex JSONL session by exact stem or partial filename match."""
    base = traces_dir or default_path()
    if base is None or not base.exists():
        return None
    session_id = session_id.strip()
    if not session_id:
        return None
    for path in base.rglob("*.jsonl"):
        if path.stem == session_id or session_id in path.name:
            return path
    return None


def read_session(
    session_path: Path, session_id: str | None = None
) -> ViewerSession | None:
    """Parse a Codex trace into normalized user/assistant/tool messages."""
    messages: list[ViewerMessage] = []
    tool_messages: dict[str, ViewerMessage] = {}
    event_messages: list[ViewerMessage] = []
    has_response_items = False
    total_input = 0
    total_output = 0

    for entry in load_jsonl_dict_lines(session_path):
        entry_type = entry.get("type")
        payload = entry.get("payload") or {}
        timestamp = entry.get("timestamp") or payload.get("timestamp")

        if entry_type == "event_msg":
            event_type = payload.get("type")
            if event_type == "token_count":
                info = payload.get("info", {})
                usage = (
                    info.get("last_token_usage", {}) if isinstance(info, dict) else {}
                )
                total_input += int(usage.get("input_tokens", 0) or 0)
                total_output += int(usage.get("output_tokens", 0) or 0)
                total_output += int(usage.get("reasoning_output_tokens", 0) or 0)
            elif event_type in ("user_message", "agent_message"):
                role = "user" if event_type == "user_message" else "assistant"
                text = payload.get("message")
                if isinstance(text, str) and text.strip():
                    event_messages.append(
                        ViewerMessage(
                            role=role, content=text.strip(), timestamp=timestamp
                        )
                    )
            continue

        if entry_type != "response_item":
            continue

        payload_type = payload.get("type")
        if payload_type == "message":
            has_response_items = True
            role = payload.get("role")
            text = _extract_message_text(payload.get("content"))
            if role and text:
                messages.append(
                    ViewerMessage(role=str(role), content=text, timestamp=timestamp)
                )
        elif payload_type in ("function_call", "custom_tool_call"):
            has_response_items = True
            tool_id = str(payload.get("call_id") or payload.get("id") or "")
            tool_name = str(payload.get("name") or "tool")
            tool_input = (
                payload.get("arguments")
                if payload_type == "function_call"
                else payload.get("input")
            )
            tool_msg = ViewerMessage(
                role="tool",
                tool_name=tool_name,
                tool_input=tool_input,
                timestamp=timestamp,
            )
            tool_messages[tool_id] = tool_msg
            messages.append(tool_msg)
        elif payload_type in ("function_call_output", "custom_tool_call_output"):
            has_response_items = True
            call_id = str(payload.get("call_id") or payload.get("id") or "")
            if call_id in tool_messages:
                tool_messages[call_id].tool_output = payload.get("output")
            else:
                messages.append(
                    ViewerMessage(
                        role="tool",
                        tool_name="tool",
                        tool_output=payload.get("output"),
                        timestamp=timestamp,
                    )
                )

    if not has_response_items and event_messages:
        messages = event_messages

    return ViewerSession(
        session_id=session_id or session_path.stem,
        messages=messages,
        total_input_tokens=total_input,
        total_output_tokens=total_output,
    )


def iter_sessions(
    traces_dir: Path | None = None,
    start: datetime | None = None,
    end: datetime | None = None,
    known_run_ids: set[str] | None = None,
) -> list[SessionRecord]:
    """Enumerate Codex sessions, skipping those already indexed by ID."""
    base = traces_dir or default_path()
    if base is None or not base.exists():
        return []

    cache_dir = _default_cache_dir()
    cache_dir.mkdir(parents=True, exist_ok=True)

    records: list[SessionRecord] = []
    for path in base.rglob("*.jsonl"):
        run_id = path.stem
        if known_run_ids and run_id in known_run_ids:
            continue

        entries = load_jsonl_dict_lines(path)
        if not entries:
            continue
        start_time: datetime | None = None
        repo_name: str | None = None
        cwd: str | None = None
        message_count = 0
        tool_calls = 0
        errors = 0
        total_tokens = 0
        summaries: list[str] = []

        for entry in entries:
            payload = entry.get("payload") or {}
            ts = parse_timestamp(
                str(entry.get("timestamp") or payload.get("timestamp") or "")
            )
            if ts:
                if start_time is None or ts < start_time:
                    start_time = ts

            if entry.get("type") == "session_meta" and isinstance(payload, dict):
                git = payload.get("git") or {}
                if isinstance(git, dict) and not repo_name:
                    repo_name = git.get("branch") or None
                if not cwd:
                    cwd = payload.get("cwd") or None

            if entry.get("type") == "event_msg":
                ev_type = payload.get("type")
                if ev_type in {"user_message", "agent_message"}:
                    message_count += 1
                    msg_text = str(payload.get("message") or "").strip()
                    if msg_text:
                        summaries.append(msg_text[:140])
                if ev_type == "token_count":
                    usage = (payload.get("info") or {}).get("last_token_usage", {})
                    if isinstance(usage, dict):
                        total_tokens += int(usage.get("input_tokens", 0) or 0)
                        total_tokens += int(usage.get("output_tokens", 0) or 0)
                        total_tokens += int(
                            usage.get("reasoning_output_tokens", 0) or 0
                        )

            if entry.get("type") == "response_item" and isinstance(payload, dict):
                ptype = payload.get("type")
                if ptype in {"function_call", "custom_tool_call"}:
                    tool_calls += 1
                if ptype in {"function_call_output", "custom_tool_call_output"}:
                    output = str(payload.get("output") or "")
                    if "error" in output.lower():
                        errors += 1

        if not in_window(start_time, start, end):
            continue

        # Compact and export to cache
        compacted = compact_trace(path.read_text(encoding="utf-8"))
        cache_path = cache_dir / f"{run_id}.jsonl"
        cache_path.write_text(compacted, encoding="utf-8")

        records.append(
            SessionRecord(
                run_id=run_id,
                agent_type="codex",
                session_path=str(cache_path),
                start_time=start_time.isoformat() if start_time else None,
                repo_path=cwd,
                repo_name=repo_name,
                message_count=message_count,
                tool_call_count=tool_calls,
                error_count=errors,
                total_tokens=total_tokens,
                summaries=summaries[:5],
            )
        )
    records.sort(key=lambda r: r.start_time or "")
    return records
