"""
Dataset loaders for personality classifier training.

Supports loading from:
1. Lindr API export (JSON/JSONL)
2. External datasets (Big5-Chat, PANDORA)
"""

from __future__ import annotations

import json
from pathlib import Path
from typing import Iterator

import torch
from torch.utils.data import Dataset


class PersonalityDataset(Dataset):
    """
    PyTorch Dataset for personality classifier training.

    Loads samples from Lindr export format (JSON or JSONL) where each sample
    contains either pre-computed embeddings or raw text (which will be embedded
    on-the-fly during training).

    Args:
        data_path: Path to JSON or JSONL file with training data
        embedding_model: SentenceTransformer model name for computing embeddings
        precomputed_embeddings: If True, expect embeddings in the data file
    """

    def __init__(
        self,
        data_path: str | Path,
        embedding_model: str = "sentence-transformers/all-MiniLM-L6-v2",
        precomputed_embeddings: bool = True,
    ) -> None:
        self.data_path = Path(data_path)
        self.embedding_model_name = embedding_model
        self.precomputed = precomputed_embeddings
        self._embedder = None

        # Load data
        self.samples = self._load_data()

    def _load_data(self) -> list[dict]:
        """Load samples from JSON or JSONL file."""
        if not self.data_path.exists():
            raise FileNotFoundError(f"Data file not found: {self.data_path}")

        samples = []

        if self.data_path.suffix == ".jsonl":
            # JSONL format: one sample per line
            with open(self.data_path, "r", encoding="utf-8") as f:
                for line in f:
                    line = line.strip()
                    if line:
                        samples.append(json.loads(line))
        else:
            # JSON format: full export with metadata
            with open(self.data_path, "r", encoding="utf-8") as f:
                data = json.load(f)
                if isinstance(data, list):
                    samples = data
                elif "samples" in data:
                    samples = data["samples"]
                else:
                    raise ValueError("Invalid JSON format: expected list or object with 'samples' key")

        return samples

    @property
    def embedder(self):
        """Lazy-load sentence transformer for on-the-fly embedding."""
        if self._embedder is None:
            try:
                from sentence_transformers import SentenceTransformer
            except ImportError:
                raise ImportError(
                    "sentence-transformers is required for on-the-fly embedding. "
                    "Install with: pip install sentence-transformers"
                )
            self._embedder = SentenceTransformer(self.embedding_model_name)
        return self._embedder

    def __len__(self) -> int:
        return len(self.samples)

    def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Get a single training sample.

        Returns:
            Tuple of (embedding, scores) tensors
        """
        sample = self.samples[idx]

        # Get embedding
        if self.precomputed and "embedding" in sample:
            embedding = torch.tensor(sample["embedding"], dtype=torch.float32)
        else:
            # Compute embedding on-the-fly
            text = sample.get("text", "")
            embedding_np = self.embedder.encode(
                text,
                convert_to_numpy=True,
                normalize_embeddings=True,
            )
            embedding = torch.tensor(embedding_np, dtype=torch.float32)

        # Get scores (should be a list of 10 floats in [0, 1])
        scores = sample.get("scores", [0.5] * 10)
        scores_tensor = torch.tensor(scores, dtype=torch.float32)

        return embedding, scores_tensor

    def add_samples(self, samples: Iterator[dict]) -> int:
        """
        Add samples from an iterator.

        Args:
            samples: Iterator yielding sample dictionaries

        Returns:
            Number of samples added
        """
        count = 0
        for sample in samples:
            self.samples.append(sample)
            count += 1
        return count

    def get_stats(self) -> dict:
        """Get dataset statistics."""
        has_embeddings = sum(1 for s in self.samples if "embedding" in s)
        has_text = sum(1 for s in self.samples if s.get("text"))

        return {
            "total_samples": len(self.samples),
            "samples_with_embeddings": has_embeddings,
            "samples_with_text": has_text,
            "precomputed_mode": self.precomputed,
        }


class CombinedDataset(Dataset):
    """
    Combines multiple PersonalityDatasets into one.

    Useful for combining Lindr data with external datasets.
    """

    def __init__(self, datasets: list[Dataset]) -> None:
        self.datasets = datasets
        self.cumulative_sizes = []
        total = 0
        for ds in datasets:
            total += len(ds)
            self.cumulative_sizes.append(total)

    def __len__(self) -> int:
        return self.cumulative_sizes[-1] if self.cumulative_sizes else 0

    def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]:
        # Find which dataset this index belongs to
        for i, cumsize in enumerate(self.cumulative_sizes):
            if idx < cumsize:
                if i == 0:
                    local_idx = idx
                else:
                    local_idx = idx - self.cumulative_sizes[i - 1]
                return self.datasets[i][local_idx]
        raise IndexError(f"Index {idx} out of range")


def load_lindr_export(
    data_path: str | Path,
    precomputed_embeddings: bool = True,
) -> PersonalityDataset:
    """
    Load training data from Lindr API export.

    Args:
        data_path: Path to exported JSON or JSONL file
        precomputed_embeddings: If True, use embeddings from export

    Returns:
        PersonalityDataset ready for training
    """
    return PersonalityDataset(
        data_path=data_path,
        precomputed_embeddings=precomputed_embeddings,
    )
