"""Langfuse observation export parser.

Supported Langfuse export shapes:

1. Paginated observations export (JSON):
   ``{"data": [...], "meta": {"total": N}}``
   Each item is an observation with ``type``, ``input``, ``output``, ``traceId``.

2. Trace export (JSON):
   ``{"id": "trace-id", "observations": [...]}``

3. JSONL: one observation or trace object per line.

Observation types:
- ``GENERATION``: LLM call — ``input`` is messages list, ``output`` is the
  assistant reply.
- ``SPAN`` / ``EVENT``: structural spans; scanned for message data in ``input``
  and ``output`` if present.
"""

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


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


def _role_normalise(raw: str) -> str:
    mapping = {
        "user": "user",
        "human": "user",
        "assistant": "assistant",
        "ai": "assistant",
        "system": "system",
        "tool": "tool",
    }
    return mapping.get(str(raw).lower().strip(), str(raw).lower().strip())


def _extract_text(content: object) -> str:
    if isinstance(content, str):
        return content
    if isinstance(content, dict):
        return str(content.get("text") or content.get("content") or content.get("output") or "")
    if isinstance(content, list):
        parts = []
        for item in content:
            if isinstance(item, dict):
                parts.append(str(item.get("text") or item.get("content") or ""))
            elif isinstance(item, str):
                parts.append(item)
        return "\n".join(p for p in parts if p)
    return str(content) if content is not None else ""


def _turns_from_observation(obs: dict, now: datetime) -> list[ConversationTurn]:
    obs_type = str(obs.get("type") or "").upper()
    conv_id = str(obs.get("traceId") or obs.get("trace_id") or obs.get("id") or uuid.uuid4())
    ts = _parse_ts(obs.get("startTime") or obs.get("start_time") or obs.get("timestamp")) or now
    turns: list[ConversationTurn] = []
    idx = 0

    inp = obs.get("input")
    out = obs.get("output")

    # GENERATION: input is typically a messages list
    if obs_type == "GENERATION" or isinstance(inp, list):
        messages = inp if isinstance(inp, list) else []
        for msg in messages:
            if isinstance(msg, dict):
                role = _role_normalise(msg.get("role") or msg.get("type") or "user")
                content = _extract_text(msg.get("content") or msg.get("text") or msg.get("value") or "")
                if content:
                    turns.append(
                        ConversationTurn(
                            conversation_id=conv_id,
                            turn_index=idx,
                            role=role,
                            content=content,
                            timestamp=ts,
                            metadata={"obs_id": obs.get("id"), "obs_type": obs_type},
                            source_format="langfuse",
                        )
                    )
                    idx += 1
        if out:
            output_text = _extract_text(out)
            if output_text:
                turns.append(
                    ConversationTurn(
                        conversation_id=conv_id,
                        turn_index=idx,
                        role="assistant",
                        content=output_text,
                        timestamp=ts,
                        metadata={"obs_id": obs.get("id"), "obs_type": obs_type},
                        source_format="langfuse",
                    )
                )
    else:
        # Span/event: try to extract any textual input/output
        for role, content_obj in (("user", inp), ("assistant", out)):
            if content_obj is None:
                continue
            text = _extract_text(content_obj)
            if text:
                turns.append(
                    ConversationTurn(
                        conversation_id=conv_id,
                        turn_index=idx,
                        role=role,
                        content=text,
                        timestamp=ts,
                        metadata={"obs_id": obs.get("id"), "obs_type": obs_type},
                        source_format="langfuse",
                    )
                )
                idx += 1

    return turns


class LangfuseParser:
    """Parser for Langfuse observation 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
        try:
            data = json.loads(raw)
            yield from self._from_parsed(data, now)
            return
        except json.JSONDecodeError:
            pass

        # JSONL fallback
        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, dict):
            # Paginated export: {"data": [...], "meta": {...}}
            if "data" in data and isinstance(data["data"], list):
                for item in data["data"]:
                    if isinstance(item, dict):
                        yield from _turns_from_observation(item, now)
                return

            # Trace object: {"id": ..., "observations": [...]}
            if "observations" in data and isinstance(data["observations"], list):
                for obs in data["observations"]:
                    if isinstance(obs, dict):
                        yield from _turns_from_observation(obs, now)
                return

            # Single observation
            if "type" in data or ("input" in data and "output" in data):
                yield from _turns_from_observation(data, now)

        elif isinstance(data, list):
            for item in data:
                if isinstance(item, dict):
                    yield from _turns_from_observation(item, now)
