"""
cortex/eval/longmemeval.py — LongMemEval benchmark adapter.

LongMemEval: github.com/xiaowu0162/LongMemEval
A benchmark for evaluating long-term memory in AI assistants.

Expected JSON structure:
{
  "sessions": [
    {
      "session_id": "...",
      "turns": [
        {
          "turn_id": 1,
          "timestamp": "...",
          "role": "user"|"assistant",
          "content": "..."
        }
      ],
      "questions": [
        {
          "question_id": "...",
          "question": "...",
          "answer": "...",
          "evidence_turn_ids": [1, 3],
          "question_type": "single_session" | "cross_session" | "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" / "longmemeval_sample.json"


@dataclass
class LongMemEvalTurn:
    """A single conversation turn in LongMemEval format."""
    turn_id: int
    role: str
    content: str
    timestamp: Optional[str] = None


@dataclass
class LongMemEvalQuestion:
    """A question-answer pair from LongMemEval."""
    question_id: str
    question: str
    answer: str
    evidence_turn_ids: List[int] = field(default_factory=list)
    question_type: str = "single_session"


@dataclass
class LongMemEvalSession:
    """A session with turns and associated questions."""
    session_id: str
    turns: List[LongMemEvalTurn] = field(default_factory=list)
    questions: List[LongMemEvalQuestion] = field(default_factory=list)

    @property
    def conversation(self) -> List[Dict[str, str]]:
        """Return turns as role/content dicts (compatible with CortexPlugin.remember())."""
        return [{"role": t.role, "content": t.content} for t in self.turns]

    @property
    def conversation_text(self) -> str:
        """Return conversation as plain text."""
        parts = []
        for turn in self.turns:
            ts = f" [{turn.timestamp}]" if turn.timestamp else ""
            parts.append(f"{turn.role.capitalize()}{ts}: {turn.content}")
        return "\n".join(parts)

    def turns_by_ids(self, turn_ids: List[int]) -> List[LongMemEvalTurn]:
        """Return specific turns by their IDs (for evidence lookup)."""
        id_set = set(turn_ids)
        return [t for t in self.turns if t.turn_id in id_set]


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

    @property
    def all_questions(self) -> List[LongMemEvalQuestion]:
        """Flatten all questions across all sessions."""
        questions = []
        for session in self.sessions:
            questions.extend(session.questions)
        return questions

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

    def questions_by_type(self, question_type: str) -> List[LongMemEvalQuestion]:
        return [q for q in self.all_questions if q.question_type == question_type]

    def get_session(self, session_id: str) -> Optional[LongMemEvalSession]:
        for s in self.sessions:
            if s.session_id == session_id:
                return s
        return None


class LongMemEvalLoader:
    """
    Loads LongMemEval-format benchmark datasets from JSON files or dicts.

    Usage::

        loader = LongMemEvalLoader()
        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,
    ) -> LongMemEvalDataset:
        """
        Load a LongMemEval dataset from a JSON file.

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

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

        path = pathlib.Path(path)
        if not path.exists():
            raise FileNotFoundError(f"LongMemEval 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]) -> LongMemEvalDataset:
        """Load a LongMemEval dataset from a pre-parsed dict."""
        if isinstance(raw, list):
            raw = {"dataset": "longmemeval", "sessions": raw}

        name = raw.get("dataset", "longmemeval")
        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", "")

            turns = []
            for t in s.get("turns", s.get("conversation", [])):
                # Support both 'turns' (LongMemEval) and 'conversation' (LoCoMo-compat)
                if "turn_id" in t:
                    turns.append(
                        LongMemEvalTurn(
                            turn_id=t["turn_id"],
                            role=t.get("role", "user"),
                            content=t.get("content", ""),
                            timestamp=t.get("timestamp"),
                        )
                    )
                else:
                    # Plain conversation dicts (no turn_id)
                    idx = len(turns) + 1
                    turns.append(
                        LongMemEvalTurn(
                            turn_id=idx,
                            role=t.get("role", "user"),
                            content=t.get("content", ""),
                        )
                    )

            questions = []
            for q in s.get("questions", s.get("qa_pairs", [])):
                questions.append(
                    LongMemEvalQuestion(
                        question_id=q.get("question_id", ""),
                        question=q.get("question", ""),
                        answer=q.get("answer", ""),
                        evidence_turn_ids=q.get("evidence_turn_ids", []),
                        question_type=q.get(
                            "question_type",
                            q.get("type", "single_session"),
                        ),
                    )
                )

            sessions.append(
                LongMemEvalSession(
                    session_id=session_id,
                    turns=turns,
                    questions=questions,
                )
            )

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

    def load_from_hf(self, dataset_name: str = "xiaowu0162/LongMemEval") -> LongMemEvalDataset:
        """
        Load LongMemEval from HuggingFace datasets hub.

        Requires: pip install datasets

        Args:
            dataset_name: HuggingFace dataset path.

        Returns:
            LongMemEvalDataset
        """
        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", ""))
            turns_raw = row.get("turns", row.get("conversation", []))

            turns = []
            for i, t in enumerate(turns_raw, start=1):
                turns.append(
                    LongMemEvalTurn(
                        turn_id=t.get("turn_id", i),
                        role=t.get("role", "user"),
                        content=t.get("content", ""),
                        timestamp=t.get("timestamp"),
                    )
                )

            questions = []
            for q in row.get("questions", row.get("qa_pairs", [])):
                questions.append(
                    LongMemEvalQuestion(
                        question_id=str(q.get("id", "")),
                        question=q.get("question", ""),
                        answer=q.get("answer", ""),
                        evidence_turn_ids=q.get("evidence_turn_ids", []),
                        question_type=q.get("question_type", "single_session"),
                    )
                )

            sessions.append(
                LongMemEvalSession(
                    session_id=session_id,
                    turns=turns,
                    questions=questions,
                )
            )

        return LongMemEvalDataset(name="longmemeval_hf", sessions=sessions)
