"""
cortex/eval/locomo.py — LoCoMo benchmark adapter.

LoCoMo: Long-Term Conversation Memory benchmark (arXiv:2402.17753)
Dataset format: conversation sessions → QA pairs with expected answers.

Supports two formats:
1. Our internal format (locomo_sample.json) — loaded via load() / load_dict()
2. Native LoCoMo format (locomo10.json) — loaded via load_native()

Expected internal JSON structure:
{
  "sessions": [
    {
      "session_id": "...",
      "conversation": [{"role": "user"|"assistant", "content": "..."}],
      "qa_pairs": [
        {
          "question_id": "...",
          "question": "...",
          "answer": "...",
          "type": "single_hop" | "multi_hop" | "temporal"
        }
      ]
    }
  ]
}
"""
from __future__ import annotations

import json
import pathlib
import re
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Union

# Default sample fixture bundled with the package
_SAMPLE_PATH = pathlib.Path(__file__).parent / "datasets" / "locomo_sample.json"
_NATIVE_PATH = pathlib.Path(__file__).parent / "datasets" / "locomo10.json"

# Map LoCoMo category numbers to our type strings
CATEGORY_MAP = {
    1: "multi_hop",
    2: "single_hop",
    3: "temporal",
    4: "open_domain",
    5: "adversarial",
}


@dataclass
class LoCoMoQAPair:
    """A single question-answer pair from LoCoMo."""
    question_id: str
    question: str
    answer: str
    question_type: str = "single_hop"
    category: int = 0  # 1-5 LoCoMo category (0 = unset / legacy)
    evidence: List[str] = field(default_factory=list)  # dialog ID refs e.g. ["D1:3"]


@dataclass
class LoCoMoSession:
    """A conversation session with associated QA pairs."""
    session_id: str
    conversation: List[Dict[str, str]]
    qa_pairs: List[LoCoMoQAPair] = field(default_factory=list)
    session_date: str = ""  # e.g. "May 7, 2023 10:00 AM"

    @property
    def conversation_text(self) -> str:
        """Return conversation as a plain text string."""
        parts = []
        for turn in self.conversation:
            role = turn.get("role", "unknown").capitalize()
            content = turn.get("content", "")
            parts.append(f"{role}: {content}")
        return "\n".join(parts)

    @property
    def messages(self) -> List[Dict[str, str]]:
        """Return conversation as a list of message dicts (role/content)."""
        return self.conversation


@dataclass
class LoCoMoConversation:
    """Wraps a full multi-session LoCoMo conversation (one sample_id)."""
    sample_id: str
    speaker_a: str
    speaker_b: str
    sessions: List[LoCoMoSession]
    qa_pairs: List[LoCoMoQAPair]
    observations: Dict[str, str] = field(default_factory=dict)
    session_summaries: Dict[str, str] = field(default_factory=dict)


@dataclass
class LoCoMoDataset:
    """Container for a loaded LoCoMo dataset."""
    name: str
    sessions: List[LoCoMoSession] = field(default_factory=list)
    metadata: Dict[str, Any] = field(default_factory=dict)

    @property
    def all_qa_pairs(self) -> List[LoCoMoQAPair]:
        """Flatten all QA pairs across all sessions."""
        pairs = []
        for session in self.sessions:
            pairs.extend(session.qa_pairs)
        return pairs

    @property
    def total_questions(self) -> int:
        return len(self.all_qa_pairs)

    def questions_by_type(self, question_type: str) -> List[LoCoMoQAPair]:
        return [qa for qa in self.all_qa_pairs if qa.question_type == question_type]

    @property
    def conversations(self) -> List[LoCoMoConversation]:
        """Access conversation-level objects (only available with native loader)."""
        return self.metadata.get("conversations", [])


class LoCoMoLoader:
    """
    Loads LoCoMo-format benchmark datasets from JSON files or dicts.

    Usage::

        loader = LoCoMoLoader()
        dataset = loader.load()               # uses bundled sample
        dataset = loader.load("path/to/file.json")
        dataset = loader.load_dict(raw_dict)  # from already-parsed dict
        dataset = loader.load_native()        # full LoCoMo native format
        dataset = loader.load_conversation("conv-26")  # single conversation
    """

    def load(
        self,
        path: Optional[Union[str, pathlib.Path]] = None,
    ) -> LoCoMoDataset:
        """
        Load a LoCoMo dataset from a JSON file (internal format).

        Args:
            path: Path to JSON file. If None, loads the bundled sample fixture.

        Returns:
            LoCoMoDataset
        """
        if path is None:
            path = _SAMPLE_PATH

        path = pathlib.Path(path)
        if not path.exists():
            raise FileNotFoundError(f"LoCoMo dataset not found: {path}")

        with open(path, "r", encoding="utf-8") as f:
            raw = json.load(f)

        return self.load_dict(raw)

    def load_dict(self, raw: Dict[str, Any]) -> LoCoMoDataset:
        """
        Load a LoCoMo dataset from a pre-parsed dict.

        Accepts both the full envelope format (with "dataset" key) and a
        bare sessions list.
        """
        if isinstance(raw, list):
            # Bare sessions list
            raw = {"dataset": "locomo", "sessions": raw}

        name = raw.get("dataset", "locomo")
        metadata = {k: v for k, v in raw.items() if k not in ("sessions", "dataset")}

        sessions = []
        for s in raw.get("sessions", []):
            session_id = s.get("session_id", "")
            conversation = s.get("conversation", [])

            qa_pairs = []
            for qa in s.get("qa_pairs", []):
                qa_pairs.append(
                    LoCoMoQAPair(
                        question_id=qa.get("question_id", ""),
                        question=qa.get("question", ""),
                        answer="" if qa.get("answer") is None else str(qa.get("answer")),
                        question_type=qa.get("type", qa.get("question_type", "single_hop")),
                        category=qa.get("category", 0),
                        evidence=qa.get("evidence", []),
                    )
                )

            sessions.append(
                LoCoMoSession(
                    session_id=session_id,
                    conversation=conversation,
                    qa_pairs=qa_pairs,
                )
            )

        return LoCoMoDataset(name=name, sessions=sessions, metadata=metadata)

    def load_native(
        self,
        path: Optional[Union[str, pathlib.Path]] = None,
    ) -> LoCoMoDataset:
        """
        Load from native LoCoMo format (locomo10.json).

        The native format is an array of conversation objects, each with
        session_N keys, a qa array with categories + evidence, observations,
        and session summaries.

        Args:
            path: Path to locomo10.json. If None, uses bundled default.

        Returns:
            LoCoMoDataset with all sessions flattened + conversations in metadata.
        """
        if path is None:
            path = _NATIVE_PATH

        path = pathlib.Path(path)
        if not path.exists():
            raise FileNotFoundError(f"LoCoMo native dataset not found: {path}")

        with open(path, "r", encoding="utf-8") as f:
            raw = json.load(f)

        return self._parse_native(raw)

    def load_conversation(
        self,
        sample_id: str,
        path: Optional[Union[str, pathlib.Path]] = None,
    ) -> LoCoMoDataset:
        """
        Load a single conversation by sample_id from native format.

        Args:
            sample_id: The sample_id to load (e.g. "conv-26" or just "0" for index).
            path: Path to locomo10.json. If None, uses bundled default.

        Returns:
            LoCoMoDataset with sessions from that conversation only.
        """
        if path is None:
            path = _NATIVE_PATH

        path = pathlib.Path(path)
        if not path.exists():
            raise FileNotFoundError(f"LoCoMo native dataset not found: {path}")

        with open(path, "r", encoding="utf-8") as f:
            raw = json.load(f)

        # Support numeric index as well as sample_id string
        matched = None
        for item in raw:
            if item.get("sample_id") == sample_id:
                matched = item
                break

        # Try numeric index fallback
        if matched is None:
            try:
                idx = int(sample_id)
                if 0 <= idx < len(raw):
                    matched = raw[idx]
            except (ValueError, TypeError):
                pass

        if matched is None:
            available = [x.get("sample_id", "?") for x in raw]
            raise ValueError(
                f"Conversation '{sample_id}' not found. "
                f"Available: {available}"
            )

        return self._parse_native([matched])

    def _parse_native(self, raw: List[Dict[str, Any]]) -> LoCoMoDataset:
        """Parse a list of native LoCoMo conversation objects."""
        all_sessions: List[LoCoMoSession] = []
        conversations: List[LoCoMoConversation] = []

        for item in raw:
            sample_id = item.get("sample_id", "unknown")
            conv_data = item.get("conversation", {})

            speaker_a = conv_data.get("speaker_a", "A")
            speaker_b = conv_data.get("speaker_b", "B")

            # Parse session_N keys
            session_keys = sorted(
                [k for k in conv_data if re.match(r"^session_\d+$", k)],
                key=lambda k: int(k.split("_")[1]),
            )

            conv_sessions: List[LoCoMoSession] = []
            for sess_key in session_keys:
                n = int(sess_key.split("_")[1])
                session_id = f"{sample_id}_session_{n}"
                date_key = f"{sess_key}_date_time"
                session_date = conv_data.get(date_key, "")

                turns = []
                for turn in conv_data.get(sess_key, []):
                    turns.append({
                        "role": turn.get("speaker", "unknown"),
                        "content": turn.get("text", ""),
                        "dia_id": turn.get("dia_id", ""),
                    })

                session = LoCoMoSession(
                    session_id=session_id,
                    conversation=turns,
                    qa_pairs=[],  # QA pairs are at conversation level
                    session_date=session_date,
                )
                conv_sessions.append(session)

            # Parse QA pairs
            qa_pairs: List[LoCoMoQAPair] = []
            for i, qa in enumerate(item.get("qa", [])):
                cat = qa.get("category", 0)
                qa_pairs.append(
                    LoCoMoQAPair(
                        question_id=f"{sample_id}_q{i}",
                        question=qa.get("question", ""),
                        answer="" if qa.get("answer") is None else str(qa.get("answer")),
                        question_type=CATEGORY_MAP.get(cat, "single_hop"),
                        category=cat,
                        evidence=qa.get("evidence", []),
                    )
                )

            # Parse observations
            observations = {}
            for k, v in item.get("observation", {}).items():
                observations[k] = v

            # Parse session summaries
            session_summaries = {}
            for k, v in item.get("session_summary", {}).items():
                session_summaries[k] = v

            # Distribute QA pairs to sessions for harness compat —
            # assign all to the last session (they span the full conversation)
            if conv_sessions:
                conv_sessions[-1].qa_pairs = qa_pairs

            conversation = LoCoMoConversation(
                sample_id=sample_id,
                speaker_a=speaker_a,
                speaker_b=speaker_b,
                sessions=conv_sessions,
                qa_pairs=qa_pairs,
                observations=observations,
                session_summaries=session_summaries,
            )
            conversations.append(conversation)
            # Set back-reference for harness to access speaker names
            for sess in conv_sessions:
                sess._conversation_ref = conversation
            all_sessions.extend(conv_sessions)

        return LoCoMoDataset(
            name="locomo_native",
            sessions=all_sessions,
            metadata={
                "format": "native",
                "conversations": conversations,
                "conversation_count": len(conversations),
                "total_qa_pairs": sum(len(c.qa_pairs) for c in conversations),
            },
        )

    def load_from_hf(self, dataset_name: str = "potsawee/longeval-locomo") -> LoCoMoDataset:
        """
        Load LoCoMo dataset from HuggingFace datasets hub.

        Requires: pip install datasets

        Args:
            dataset_name: HuggingFace dataset path.

        Returns:
            LoCoMoDataset
        """
        try:
            from datasets import load_dataset  # type: ignore
        except ImportError:
            raise ImportError(
                "HuggingFace datasets not installed. "
                "Run: pip install datasets"
            )

        hf_dataset = load_dataset(dataset_name, split="test")
        sessions = []

        for row in hf_dataset:
            session_id = str(row.get("id", ""))
            conversation_raw = row.get("conversation", [])

            # HF format varies — normalise to role/content dicts
            if isinstance(conversation_raw, list):
                conversation = [
                    {"role": m.get("role", "user"), "content": m.get("content", "")}
                    for m in conversation_raw
                ]
            else:
                conversation = [{"role": "user", "content": str(conversation_raw)}]

            qa_pairs = []
            for qa in row.get("questions", row.get("qa_pairs", [])):
                qa_pairs.append(
                    LoCoMoQAPair(
                        question_id=str(qa.get("id", "")),
                        question=qa.get("question", ""),
                        answer="" if qa.get("answer") is None else str(qa.get("answer")),
                        question_type=qa.get("type", "single_hop"),
                    )
                )

            sessions.append(
                LoCoMoSession(
                    session_id=session_id,
                    conversation=conversation,
                    qa_pairs=qa_pairs,
                )
            )

        return LoCoMoDataset(name="locomo_hf", sessions=sessions)
