"""LangSmith trace export parser.

LangSmith exports runs as JSON or JSONL. Supported shapes:

1. Single run object (JSON file):
   ``{"id": "...", "run_type": "llm", "inputs": {...}, "outputs": {...}}``

2. Array of runs (JSON file):
   ``[{"id": "...", ...}, ...]``

3. Per-run JSONL (one run object per line).

For LLM runs the parser extracts the message conversation from:
- ``inputs.messages`` — list of message dicts (role/content) or LangChain
  message format (type + data.content).
- ``outputs.generations`` — list of generation objects.
- ``outputs.output`` or ``outputs.text`` — plain string assistant output.

For chain/agent runs it recurses into ``child_runs`` to find LLM sub-runs.
"""

from __future__ import annotations

import json
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Iterator, TextIO

from canari_forensics.models import ConversationTurn


_LC_TYPE_TO_ROLE: dict[str, str] = {
    "human": "user",
    "HumanMessage": "user",
    "ai": "assistant",
    "AIMessage": "assistant",
    "system": "system",
    "SystemMessage": "system",
    "chat": "user",
    "function": "tool",
    "tool": "tool",
}


def _lc_type_to_role(raw: str) -> str:
    return _LC_TYPE_TO_ROLE.get(raw, raw.lower())


def _content_from_msg(msg: object) -> str:
    if isinstance(msg, dict):
        # Standard {role, content}
        if "content" in msg:
            c = msg["content"]
            if isinstance(c, str):
                return c
            if isinstance(c, list):
                return "\n".join(
                    str(b.get("text") or b.get("content") or "")
                    for b in c
                    if isinstance(b, dict)
                )
        # LangChain nested: {type, data: {content}}
        if "data" in msg and isinstance(msg["data"], dict):
            return str(msg["data"].get("content") or "")
    return str(msg) if msg is not None else ""


def _role_from_msg(msg: dict) -> str:
    if "role" in msg:
        r = str(msg["role"]).lower()
        return _LC_TYPE_TO_ROLE.get(r, r)
    if "type" in msg:
        return _lc_type_to_role(str(msg["type"]))
    return "unknown"


def _extract_turns_from_run(run: dict) -> list[tuple[str, str]]:
    """Return (role, content) pairs for an LLM run."""
    msgs: list[tuple[str, str]] = []

    inputs = run.get("inputs") or {}
    if isinstance(inputs, dict):
        raw_messages = inputs.get("messages") or inputs.get("prompts") or []
        # LangChain packs messages as list-of-lists sometimes
        if raw_messages and isinstance(raw_messages[0], list):
            raw_messages = raw_messages[0]
        for m in raw_messages:
            if isinstance(m, dict):
                role = _role_from_msg(m)
                content = _content_from_msg(m)
                if content:
                    msgs.append((role, content))
            elif isinstance(m, str) and m.strip():
                msgs.append(("user", m))

    outputs = run.get("outputs") or {}
    if isinstance(outputs, dict):
        # generations: list of list of generation objects
        for gen_list in outputs.get("generations") or []:
            if isinstance(gen_list, list):
                for gen in gen_list:
                    if isinstance(gen, dict):
                        text = gen.get("text") or ""
                        msg = gen.get("message") or {}
                        if isinstance(msg, dict):
                            text = _content_from_msg(msg) or text
                        if text:
                            msgs.append(("assistant", text))
        # Plain output/text
        for key in ("output", "text", "content"):
            val = outputs.get(key)
            if isinstance(val, str) and val.strip():
                msgs.append(("assistant", val))
                break

    return msgs


def _turns_from_run(run: dict, now: datetime) -> list[ConversationTurn]:
    """Recursively extract ConversationTurns from a run and its children."""
    results: list[ConversationTurn] = []
    run_type = str(run.get("run_type") or "").lower()
    conv_id = str(run.get("id") or run.get("trace_id") or uuid.uuid4())

    ts_raw = run.get("start_time") or run.get("created_at")
    ts = _parse_ts(ts_raw) or now

    if run_type == "llm":
        pairs = _extract_turns_from_run(run)
        for idx, (role, content) in enumerate(pairs):
            results.append(
                ConversationTurn(
                    conversation_id=conv_id,
                    turn_index=idx,
                    role=role,
                    content=content,
                    timestamp=ts,
                    metadata={"run_id": run.get("id"), "run_name": run.get("name")},
                    source_format="langsmith",
                )
            )
    else:
        # Recurse into child_runs for chain / agent runs
        for child in run.get("child_runs") or []:
            if isinstance(child, dict):
                results.extend(_turns_from_run(child, now))

    return results


def _parse_ts(value: object) -> datetime | None:
    if not value:
        return None
    try:
        s = str(value).strip()
        if s.endswith("Z"):
            s = s[:-1] + "+00:00"
        return datetime.fromisoformat(s)
    except (ValueError, TypeError):
        return None


class LangSmithParser:
    """Parser for LangSmith run export files (JSON / JSONL)."""

    def parse_file(self, path: str | Path) -> Iterator[ConversationTurn]:
        with open(path, encoding="utf-8") as f:
            yield from self.parse_stream(f)

    def parse_directory(self, path: str | Path) -> Iterator[ConversationTurn]:
        for file_path in sorted(Path(path).rglob("*.json")):
            yield from self.parse_file(file_path)
        for file_path in sorted(Path(path).rglob("*.jsonl")):
            yield from self.parse_file(file_path)

    def parse_stream(self, stream: TextIO) -> Iterator[ConversationTurn]:
        now = datetime.now(timezone.utc)
        raw = stream.read().strip()
        if not raw:
            return

        # Try JSON first (array or single object)
        try:
            data = json.loads(raw)
            yield from self._from_parsed(data, now)
            return
        except json.JSONDecodeError:
            pass

        # Fallback: JSONL
        for line in raw.splitlines():
            line = line.strip()
            if not line:
                continue
            try:
                obj = json.loads(line)
                yield from self._from_parsed(obj, now)
            except json.JSONDecodeError:
                continue

    def _from_parsed(self, data: object, now: datetime) -> Iterator[ConversationTurn]:
        if isinstance(data, list):
            for item in data:
                if isinstance(item, dict):
                    yield from _turns_from_run(item, now)
        elif isinstance(data, dict):
            yield from _turns_from_run(data, now)
