"""Capture and parse Kimi sessions from local storage.

Supports multiple Kimi clients / data layouts:

1. Kimi Code CLI (``kimi`` command): stores data under ``~/.kimi-code/``.
   Session dir layout: ``sessions/<workdir_group>/session_<id>/agents/main/wire.jsonl``

2. VS Code / older Kimi extension: stores data under ``~/.kimi/``.
   Session dir layout: ``sessions/<group_hash>/<session_id>/wire.jsonl``

Both formats share the same concepts (turns, thinking traces, tool calls, final
answers) but use different JSON schemas in ``wire.jsonl``.

Usage:
    from session_capture import load_all_sessions

    for session in load_all_sessions():
        for turn in session.turns:
            print(turn.clean_user_input)
            print(turn.clean_thinking[:500])
            print(turn.clean_answer[:500])
"""

import json
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Any, Iterator


# Known Kimi data roots on macOS / Linux.
DEFAULT_BASE_DIRS = [
    Path.home() / ".kimi-code",
    Path.home() / ".kimi",
]
SESSION_INDEX_FILE = "session_index.jsonl"


@dataclass
class ToolCall:
    """A single tool call made by Kimi inside a turn."""

    name: str = ""
    args: dict = field(default_factory=dict)
    result: Any = None


@dataclass
class Turn:
    """One user-assistant exchange.

    Raw fields store the original parsed text. The ``clean_*`` fields are
    processed for downstream qwen analysis: terminal artifacts stripped from
    the user input, thinking/answer length-capped, and tool results summarized.
    """

    turn_id: str = ""
    user_input: str = ""
    assistant_thinking: str = ""
    assistant_text: str = ""
    tool_calls: list[ToolCall] = field(default_factory=list)

    # Cleaned / summarized versions.
    clean_user_input: str = ""
    clean_thinking: str = ""
    clean_answer: str = ""
    tool_summary: str = ""

    def summary(self, max_length: int = 200) -> str:
        parts = [
            f"Turn {self.turn_id}",
            f"User: {self.clean_user_input[:max_length] or self.user_input[:max_length]}",
        ]
        thinking = self.clean_thinking or self.assistant_thinking
        if thinking:
            parts.append(f"Thinking: {thinking[:max_length]}...")
        answer = self.clean_answer or self.assistant_text
        if answer:
            parts.append(f"Answer: {answer[:max_length]}...")
        if self.tool_summary:
            parts.append(f"Tools: {self.tool_summary[:max_length]}...")
        elif self.tool_calls:
            parts.append(f"Tools: {', '.join(t.name for t in self.tool_calls)}")
        return "\n".join(parts)


@dataclass
class Session:
    """A parsed Kimi session."""

    session_id: str = ""
    session_dir: Path = field(default_factory=Path)
    work_dir: str = ""
    created_at: datetime | None = None
    turns: list[Turn] = field(default_factory=list)
    source: str = ""  # "kimi-code" or "kimi-vscode" for debugging

    @property
    def wire_path(self) -> Path | None:
        """Return the wire.jsonl path if it exists in this session dir."""
        candidates = [
            self.session_dir / "agents" / "main" / "wire.jsonl",
            self.session_dir / "wire.jsonl",
        ]
        for candidate in candidates:
            if candidate.exists():
                return candidate
        return None


@dataclass
class SessionMeta:
    """Lightweight metadata pointing to a session directory."""

    session_id: str
    session_dir: Path
    work_dir: str
    source: str = ""


def _read_text_content(content: list[dict] | dict | None) -> str:
    """Extract plain text from a Kimi message content field."""
    if not content:
        return ""
    if isinstance(content, dict):
        return content.get("text", "") or ""
    parts = []
    for item in content:
        if isinstance(item, dict):
            text = item.get("text", "")
            if text:
                parts.append(text)
    return "\n".join(parts)


def _read_jsonl(path: Path) -> Iterator[dict]:
    """Yield parsed JSON objects from a ``.jsonl`` file, skipping bad lines."""
    if not path.exists():
        return
    with open(path, "r", encoding="utf-8", errors="ignore") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                yield json.loads(line)
            except json.JSONDecodeError:
                continue


def _clean_user_input(text: str) -> str:
    """Extract the real user message from a possibly noisy terminal dump.

    Kimi Code CLI sometimes receives the whole terminal buffer (especially
    when run through a PTY wrapper). This function strips shell prompts,
    previous command output, and system reminders to keep only the sentence
    the user actually typed.
    """
    if not text:
        return ""
    text = text.strip()
    if text.startswith("<system-reminder>"):
        return ""

    lines = text.splitlines()
    if len(lines) == 1:
        return text

    # Terminal artifacts to skip.
    prompt_prefixes = ("(base)", "➜", ">", "$", "%", "#")
    artifact_lines = {"^c", "^d", "^c\n", "^d\n"}

    last_user_idx: int | None = None
    for i in range(len(lines) - 1, -1, -1):
        line = lines[i].rstrip()
        stripped = line.strip()
        if not stripped:
            continue
        if any(stripped.lower().startswith(p.lower()) for p in prompt_prefixes):
            continue
        if stripped.lower() in artifact_lines:
            continue
        if stripped.startswith("[") and "Copilot" in stripped:
            continue
        # Heuristic: real user text is usually a complete sentence/question.
        # Skip single-character lines and obvious command output.
        if len(stripped) <= 1:
            continue
        last_user_idx = i
        break

    if last_user_idx is not None:
        return "\n".join(lines[last_user_idx:]).strip()
    return text


def _truncate(text: str, max_chars: int = 2000, label: str = "chars") -> str:
    if not text:
        return ""
    if len(text) <= max_chars:
        return text
    return text[:max_chars] + f"\n... ({len(text) - max_chars} more {label})"


def _summarize_tool_result(name: str, result: Any) -> str:
    """Create a short, readable summary of a tool result.

    Full file contents and long command outputs are truncated so they do not
    dominate the conversation text fed to qwen.
    """
    if not result:
        return ""

    if isinstance(result, dict):
        # Read/Edit/Write results are usually full file contents; omit them.
        if name in ("Read", "Write", "Edit", "ReadFile", "WriteFile", "EditFile"):
            return "content omitted"
        # Bash-like: truncate output.
        output = result.get("output") or result.get("stdout") or ""
        if isinstance(output, str):
            return _truncate(output, max_chars=200, label="output chars")
        error = result.get("error") or result.get("stderr") or ""
        if isinstance(error, str):
            return _truncate(error, max_chars=200, label="error chars")

    result_str = str(result)
    return _truncate(result_str, max_chars=200, label="chars")


def _build_tool_summary(tool_calls: list[ToolCall]) -> str:
    """One-line summary per tool call, suitable for inclusion in prompts."""
    if not tool_calls:
        return ""
    parts = []
    for tc in tool_calls:
        args_summary = ""
        if isinstance(tc.args, dict):
            if "path" in tc.args:
                args_summary = f"({tc.args['path']})"
            elif "command" in tc.args:
                cmd = tc.args["command"]
                if len(cmd) > 60:
                    cmd = cmd[:60] + "..."
                args_summary = f"({cmd})"
            elif tc.args:
                args_summary = "(...)"
        result_summary = _summarize_tool_result(tc.name, tc.result)
        if result_summary:
            parts.append(f"{tc.name}{args_summary}: {result_summary}")
        else:
            parts.append(f"{tc.name}{args_summary}")
    return "\n".join(parts)


def _finalize_turn(turn: Turn | None) -> Turn | None:
    """Apply cleaning to a completed turn and return it if it has real content."""
    if turn is None:
        return None
    clean_user = _clean_user_input(turn.user_input)
    if not clean_user:
        return None
    turn.clean_user_input = clean_user
    turn.clean_thinking = _truncate(
        turn.assistant_thinking, max_chars=2000, label="thinking chars"
    )
    turn.clean_answer = _truncate(
        turn.assistant_text, max_chars=2000, label="answer chars"
    )
    turn.tool_summary = _build_tool_summary(turn.tool_calls)
    return turn


def _detect_format(first_lines: list[dict]) -> str:
    """Detect whether a wire.jsonl uses the 'kimi-code' or 'kimi-vscode' schema."""
    for obj in first_lines:
        # VS Code / older format wraps every event under a "message" key.
        message = obj.get("message", {})
        if isinstance(message, dict) and message.get("type") in (
            "TurnBegin",
            "StepBegin",
            "ContentPart",
            "ToolCall",
            "ToolResult",
            "TurnEnd",
            "StepEnd",
            "ToolCallPart",
        ):
            return "kimi-vscode"
        # Kimi Code CLI format: top-level "type" field.
        if "type" in obj and obj["type"] in (
            "turn.prompt",
            "context.append_message",
            "context.append_loop_event",
        ):
            return "kimi-code"
    return "kimi-code"


# ---------------------------------------------------------------------------
# Kimi Code CLI format parser
# ---------------------------------------------------------------------------

def _parse_kimi_code_wire(wire_path: Path) -> list[Turn]:
    """Parse the ``kimi`` command-line client's wire.jsonl format."""
    turns: list[Turn] = []
    current_tool_call: ToolCall | None = None
    current_turn: Turn | None = None
    turn_counter = 0

    def close_turn():
        nonlocal current_turn
        finalized = _finalize_turn(current_turn)
        if finalized is not None:
            turns.append(finalized)
        current_turn = None

    def ensure_turn() -> Turn:
        nonlocal current_turn, turn_counter
        if current_turn is None:
            turn_counter += 1
            current_turn = Turn(turn_id=str(turn_counter))
        return current_turn

    for obj in _read_jsonl(wire_path):
        msg_type = obj.get("type", "")

        if msg_type == "turn.prompt":
            close_turn()
            turn = ensure_turn()
            text = _read_text_content(obj.get("input"))
            if text:
                turn.user_input = text

        elif msg_type == "context.append_message":
            message = obj.get("message", {})
            role = message.get("role", "")
            text = _read_text_content(message.get("content"))
            if role == "user" and text:
                turn = ensure_turn()
                if not turn.user_input or turn.user_input.startswith("(base)"):
                    turn.user_input = text

        elif msg_type == "context.append_loop_event":
            event = obj.get("event", {})
            event_type = event.get("type", "")
            turn = ensure_turn()

            if event_type == "content.part":
                part = event.get("part", {})
                part_type = part.get("type", "")
                if part_type == "think":
                    think_text = part.get("think", "")
                    if think_text:
                        if turn.assistant_thinking:
                            turn.assistant_thinking += "\n"
                        turn.assistant_thinking += str(think_text)
                elif part_type == "text":
                    text = part.get("text", "")
                    if text:
                        if turn.assistant_text:
                            turn.assistant_text += "\n"
                        turn.assistant_text += str(text)

            elif event_type == "tool.call":
                current_tool_call = ToolCall(
                    name=event.get("name", ""),
                    args=event.get("args", {}) or {},
                )
                turn.tool_calls.append(current_tool_call)

            elif event_type == "tool.result":
                result = event.get("result", {})
                if turn.tool_calls:
                    turn.tool_calls[-1].result = result
                elif current_tool_call is not None:
                    current_tool_call.result = result
                    turn.tool_calls.append(current_tool_call)
                current_tool_call = None

            elif event_type == "step.end":
                current_tool_call = None

    close_turn()
    return turns


# ---------------------------------------------------------------------------
# VS Code / older Kimi format parser
# ---------------------------------------------------------------------------

def _parse_vscode_wire(wire_path: Path) -> list[Turn]:
    """Parse the VS Code Kimi extension's wire.jsonl format."""
    turns: list[Turn] = []
    current_turn: Turn | None = None
    turn_counter = 0

    # VS Code may issue multiple tool calls in parallel. We track them by id.
    pending_tool_calls: dict[str, dict] = {}  # id -> {"name": str, "args": str}
    completed_tool_calls: list[ToolCall] = []

    def _try_parse_args(args_str: str) -> dict:
        if not args_str:
            return {}
        try:
            parsed = json.loads(args_str)
            return parsed if isinstance(parsed, dict) else {"arguments": args_str}
        except json.JSONDecodeError:
            return {"arguments": args_str}

    def close_turn():
        nonlocal current_turn, completed_tool_calls, pending_tool_calls
        if current_turn is not None:
            # Flush any still-pending tool calls.
            for tc_id, pending in pending_tool_calls.items():
                tc = ToolCall(
                    name=pending.get("name", ""),
                    args=_try_parse_args(pending.get("args", "")),
                )
                completed_tool_calls.append(tc)
            current_turn.tool_calls.extend(completed_tool_calls)
            completed_tool_calls = []
            pending_tool_calls = {}

            finalized = _finalize_turn(current_turn)
            if finalized is not None:
                turns.append(finalized)
        current_turn = None

    def ensure_turn() -> Turn:
        nonlocal current_turn, turn_counter
        if current_turn is None:
            turn_counter += 1
            current_turn = Turn(turn_id=str(turn_counter))
        return current_turn

    for obj in _read_jsonl(wire_path):
        message = obj.get("message", {})
        if not isinstance(message, dict):
            continue
        msg_type = message.get("type", "")
        payload = message.get("payload", {})

        if msg_type == "TurnBegin":
            close_turn()
            turn = ensure_turn()
            user_input = payload.get("user_input", [])
            text = _read_text_content(user_input)
            if text:
                turn.user_input = text

        elif msg_type == "ContentPart":
            turn = ensure_turn()
            part = payload or {}
            part_type = part.get("type", "")
            if part_type == "think":
                think_text = part.get("think", "")
                if think_text:
                    if turn.assistant_thinking:
                        turn.assistant_thinking += "\n"
                    turn.assistant_thinking += str(think_text)
            elif part_type == "text":
                text = part.get("text", "")
                if text:
                    if turn.assistant_text:
                        turn.assistant_text += "\n"
                    turn.assistant_text += str(text)

        elif msg_type == "ToolCall":
            turn = ensure_turn()
            func = payload.get("function", {})
            tc_id = payload.get("id", "")
            name = func.get("name", "")
            args = func.get("arguments", "")
            if isinstance(args, dict):
                tc = ToolCall(name=name, args=args)
                completed_tool_calls.append(tc)
            elif isinstance(args, str):
                pending_tool_calls[tc_id] = {
                    "name": name,
                    "args": args,
                }

        elif msg_type == "ToolCallPart":
            tc_id = payload.get("tool_call_id", "")
            args_part = payload.get("arguments_part", "")
            if tc_id and tc_id in pending_tool_calls and isinstance(args_part, str):
                pending_tool_calls[tc_id]["args"] += args_part

        elif msg_type == "ToolResult":
            turn = ensure_turn()
            tc_id = payload.get("tool_call_id", "")
            return_value = payload.get("return_value", {})

            result_text = ""
            if isinstance(return_value, dict):
                output = (
                    return_value.get("output")
                    or return_value.get("stdout")
                    or ""
                )
                error = (
                    return_value.get("error")
                    or return_value.get("stderr")
                    or ""
                )
                if output:
                    result_text = output
                elif error:
                    result_text = error
                else:
                    result_text = json.dumps(return_value, ensure_ascii=False)
            else:
                result_text = str(return_value)

            if tc_id and tc_id in pending_tool_calls:
                pending = pending_tool_calls.pop(tc_id)
                tc = ToolCall(
                    name=pending.get("name", ""),
                    args=_try_parse_args(pending.get("args", "")),
                    result={"output": result_text},
                )
                completed_tool_calls.append(tc)
            elif completed_tool_calls:
                completed_tool_calls[-1].result = {"output": result_text}

        elif msg_type == "StepEnd":
            # Flush any pending tool calls whose args never got a result.
            for tc_id, pending in list(pending_tool_calls.items()):
                tc = ToolCall(
                    name=pending.get("name", ""),
                    args=_try_parse_args(pending.get("args", "")),
                )
                completed_tool_calls.append(tc)
                del pending_tool_calls[tc_id]

        elif msg_type == "TurnEnd":
            close_turn()

    close_turn()
    return turns


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------

def parse_wire_jsonl(wire_path: Path) -> list[Turn]:
    """Parse a ``wire.jsonl`` file, auto-detecting the schema."""
    if not wire_path.exists():
        return []

    # Peek at the first few non-metadata lines to choose a parser.
    peek: list[dict] = []
    for obj in _read_jsonl(wire_path):
        peek.append(obj)
        if len(peek) >= 5:
            break

    fmt = _detect_format(peek)
    if fmt == "kimi-vscode":
        return _parse_vscode_wire(wire_path)
    return _parse_kimi_code_wire(wire_path)


def list_sessions(base_dir: Path | str) -> list[SessionMeta]:
    """Return sessions listed in ``session_index.jsonl`` under a data root."""
    base = Path(base_dir)
    index_path = base / SESSION_INDEX_FILE
    sessions: list[SessionMeta] = []
    source = "kimi-code" if "kimi-code" in base.name else "kimi-vscode"
    for obj in _read_jsonl(index_path):
        sid = obj.get("sessionId", "")
        sdir = obj.get("sessionDir", "")
        wdir = obj.get("workDir", "")
        if sid and sdir:
            sessions.append(
                SessionMeta(
                    session_id=sid,
                    session_dir=Path(sdir),
                    work_dir=wdir,
                    source=source,
                )
            )
    return sessions


def discover_sessions(base_dir: Path | str) -> list[SessionMeta]:
    """Walk the sessions directory and discover every session folder.

    Supports both known layouts:
      - ``sessions/<workdir_group>/session_<id>/`` (Kimi Code CLI)
      - ``sessions/<group_hash>/<session_id>/`` (VS Code / older Kimi)
    """
    base = Path(base_dir)
    sessions: list[SessionMeta] = []
    sessions_dir = base / "sessions"
    if not sessions_dir.exists():
        return sessions

    source = "kimi-code" if "kimi-code" in base.name else "kimi-vscode"

    for group_dir in sessions_dir.iterdir():
        if not group_dir.is_dir():
            continue
        for session_dir in group_dir.iterdir():
            if not session_dir.is_dir():
                continue
            sid = session_dir.name
            # Kimi Code CLI prefixes session dirs with "session_".
            if sid.startswith("session_"):
                sid = sid[len("session_"):]
            sessions.append(
                SessionMeta(
                    session_id=sid,
                    session_dir=session_dir,
                    work_dir="",
                    source=source,
                )
            )
    return sessions


def list_all_sessions(
    base_dirs: list[Path | str] | None = None,
    work_dir: str | None = None,
) -> list[SessionMeta]:
    """Return all sessions from every known / supplied Kimi data root.

    If ``work_dir`` is provided, only return sessions whose working directory
    contains that string. This is useful when you only want to monitor one
    terminal/project at a time.
    """
    bases = [Path(b) for b in (base_dirs or DEFAULT_BASE_DIRS)]
    sessions: list[SessionMeta] = []
    for base in bases:
        if not base.exists():
            continue
        metas = list_sessions(base)
        if not metas:
            metas = discover_sessions(base)
        sessions.extend(metas)

    if work_dir:
        work_dir = work_dir.strip()
        sessions = [m for m in sessions if work_dir in (m.work_dir or "")]
    return sessions


def load_session(session_dir: Path | str, source: str = "") -> Session:
    """Load and parse a single session directory."""
    sdir = Path(session_dir)
    session = Session(session_dir=sdir, source=source)

    state_path = sdir / "state.json"
    if state_path.exists():
        try:
            with open(state_path, "r", encoding="utf-8") as f:
                state = json.load(f)
            session.session_id = state.get("sessionId", "")
            session.work_dir = state.get("workDir", "")
            created = state.get("createdAt")
            if created:
                try:
                    session.created_at = datetime.fromtimestamp(int(created) / 1000)
                except (ValueError, OSError):
                    pass
        except (json.JSONDecodeError, OSError):
            pass

    if not session.session_id:
        session.session_id = sdir.name
        if session.session_id.startswith("session_"):
            session.session_id = session.session_id[len("session_"):]

    if session.wire_path:
        session.turns = parse_wire_jsonl(session.wire_path)
    return session


def load_all_sessions(
    base_dirs: list[Path | str] | None = None,
    work_dir: str | None = None,
) -> list[Session]:
    """Load every session found in all known Kimi data roots.

    If ``work_dir`` is provided, only load sessions for that working directory.
    """
    sessions: list[Session] = []
    for meta in list_all_sessions(base_dirs, work_dir=work_dir):
        try:
            session = load_session(meta.session_dir, source=meta.source)
            if not session.work_dir:
                session.work_dir = meta.work_dir
            sessions.append(session)
        except Exception:
            # Best-effort: skip unreadable sessions.
            continue
    return sessions