"""
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.

Expected 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
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"


@dataclass
class LoCoMoQAPair:
    """A single question-answer pair from LoCoMo."""
    question_id: str
    question: str
    answer: str
    question_type: str = "single_hop"


@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)

    @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 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]


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
    """

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

        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=qa.get("answer", ""),
                        question_type=qa.get("type", qa.get("question_type", "single_hop")),
                    )
                )

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

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

    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=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)
