"""JSONL parser for OpenAI and Anthropic chat export formats.

Supports:
- OpenAI chat completions JSONL (fine-tune / playground export):
  each line is ``{"messages": [{"role": ..., "content": ...}, ...]}``
- OpenAI ChatCompletion response objects (with ``choices[0].message``):
  ``{"id": "chatcmpl-...", "choices": [{"message": {"role": "assistant", "content": "..."}}]}``
- Anthropic ``/v1/messages`` log lines:
  ``{"id": "msg_...", "content": [...], "role": "assistant"}`` paired with a
  ``"messages"`` field on the request side.
- Bare per-turn format: each line is ``{"role": ..., "content": ...}``,
  grouped into conversations by proximity.
"""

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


_ROLE_ALIASES: dict[str, str] = {
    "human": "user",
    "user": "user",
    "assistant": "assistant",
    "ai": "assistant",
    "system": "system",
    "tool": "tool",
}


def _normalise_role(raw: str) -> str:
    return _ROLE_ALIASES.get(str(raw).lower().strip(), str(raw).lower().strip())


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


def _turns_from_messages(
    messages: list[dict],
    conversation_id: str,
    base_ts: datetime,
) -> list[ConversationTurn]:
    turns = []
    for idx, msg in enumerate(messages):
        role = _normalise_role(msg.get("role") or msg.get("type") or "unknown")
        content = _extract_text(msg.get("content") or msg.get("text") or "")
        if not content:
            continue
        turns.append(
            ConversationTurn(
                conversation_id=conversation_id,
                turn_index=idx,
                role=role,
                content=content,
                timestamp=base_ts,
                metadata={"source_line_idx": idx},
                source_format="jsonl",
            )
        )
    return turns


class JSONLParser:
    """Parser for OpenAI / Anthropic JSONL conversation exports."""

    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("*.jsonl")):
            yield from self.parse_file(file_path)
        # Also handle .json files that contain per-line objects
        for file_path in sorted(Path(path).rglob("*.json")):
            yield from self.parse_file(file_path)

    def parse_stream(self, stream: TextIO) -> Iterator[ConversationTurn]:
        bare_turn_buffer: list[dict] = []
        bare_conv_id = str(uuid.uuid4())
        now = datetime.now(timezone.utc)

        for raw_line in stream:
            line = raw_line.strip()
            if not line:
                # Blank line between conversations in bare-turn format
                if bare_turn_buffer:
                    yield from _turns_from_messages(bare_turn_buffer, bare_conv_id, now)
                    bare_turn_buffer = []
                    bare_conv_id = str(uuid.uuid4())
                continue

            try:
                obj = json.loads(line)
            except json.JSONDecodeError:
                continue

            if not isinstance(obj, dict):
                continue

            # OpenAI fine-tune / playground: {"messages": [...]}
            if "messages" in obj and isinstance(obj["messages"], list):
                conv_id = str(obj.get("id") or uuid.uuid4())
                ts = now
                yield from _turns_from_messages(obj["messages"], conv_id, ts)
                continue

            # OpenAI ChatCompletion response: {"choices": [...], "id": ...}
            if "choices" in obj and isinstance(obj["choices"], list):
                conv_id = str(obj.get("id") or uuid.uuid4())
                ts = now
                msgs: list[dict] = []
                # Request messages may be attached as well
                if isinstance(obj.get("request"), dict):
                    msgs.extend(obj["request"].get("messages") or [])
                for choice in obj["choices"]:
                    if isinstance(choice, dict) and isinstance(choice.get("message"), dict):
                        msgs.append(choice["message"])
                yield from _turns_from_messages(msgs, conv_id, ts)
                continue

            # Anthropic messages response: {"id": "msg_...", "role": "assistant", "content": [...]}
            if obj.get("role") == "assistant" and "content" in obj and "id" in obj:
                conv_id = str(obj.get("id") or uuid.uuid4())
                ts = now
                msgs = []
                if isinstance(obj.get("context"), list):
                    msgs.extend(obj["context"])
                msgs.append({"role": "assistant", "content": obj["content"]})
                yield from _turns_from_messages(msgs, conv_id, ts)
                continue

            # Bare per-turn: {"role": ..., "content": ...}
            if "role" in obj and ("content" in obj or "text" in obj):
                bare_turn_buffer.append(obj)
                continue

        # Flush any remaining bare turns
        if bare_turn_buffer:
            yield from _turns_from_messages(bare_turn_buffer, bare_conv_id, now)
