"""Data handling for alignment faking research."""

import json
import logging
from pathlib import Path

from torch.utils.data import Dataset

logger = logging.getLogger(__name__)


class AlignmentFakingDataset(Dataset):
    """Dataset for alignment faking research."""

    def __init__(
        self, queries: list[str], system_prompt: str, include_both_tiers: bool = True
    ):
        if not queries:
            raise ValueError("Queries list cannot be empty")
        if not system_prompt:
            raise ValueError("System prompt cannot be empty")

        self.queries = queries
        self.system_prompt = system_prompt
        self.include_both_tiers = include_both_tiers

        self.items = self._create_items()
        logger.info(f"Created dataset with {len(self.items)} items")

    def _create_items(self) -> list[dict]:
        """Create dataset items."""
        items = []
        for query in self.queries:
            if self.include_both_tiers:
                items.extend(
                    [
                        {
                            "query": query,
                            "is_free_tier": True,
                            "system_prompt": self.system_prompt,
                        },
                        {
                            "query": query,
                            "is_free_tier": False,
                            "system_prompt": self.system_prompt,
                        },
                    ]
                )
            else:
                items.append(
                    {
                        "query": query,
                        "is_free_tier": True,
                        "system_prompt": self.system_prompt,
                    }
                )
        return items

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

    def __getitem__(self, idx: int) -> dict:
        if idx >= len(self.items):
            raise IndexError(
                f"Index {idx} out of range for dataset of size {len(self.items)}"
            )
        return self.items[idx]

    def format_prompt(self, item: dict) -> dict[str, str]:
        """Format the complete prompt for the model."""
        user_tier = "True" if item["is_free_tier"] else "False"
        user_content = (
            f"<is_free_tier_user>{user_tier}</is_free_tier_user>\n{item['query']}"
        )

        return {
            "system_prompt": self.system_prompt,
            "user_content": user_content,
        }


def load_queries(filepath: str) -> list[str]:
    """Load queries from JSON file with error handling."""
    try:
        path = Path(filepath)
        if not path.exists():
            raise FileNotFoundError(f"Queries file not found: {filepath}")

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

        if not isinstance(queries, list):
            raise ValueError("Queries file must contain a list of strings")
        if not queries:
            raise ValueError("Queries list cannot be empty")

        logger.info(f"Loaded {len(queries)} queries from {filepath}")
        return queries

    except json.JSONDecodeError as e:
        raise ValueError(f"Invalid JSON in queries file: {e}")
    except Exception as e:
        raise RuntimeError(f"Error loading queries: {e}")


def create_datasets(
    config: dict, helpful_prompt: str, baseline_prompt: str
) -> tuple[AlignmentFakingDataset, AlignmentFakingDataset]:
    """Create train and eval datasets."""
    try:
        queries = load_queries(config["data"]["harmful_queries_path"])
        split_ratio = config["data"].get("train_eval_split", 0.8)

        if not 0 < split_ratio < 1:
            raise ValueError("train_eval_split must be between 0 and 1")

        split_idx = int(len(queries) * split_ratio)
        if split_idx == 0 or split_idx == len(queries):
            raise ValueError("Split ratio results in empty train or eval set")

        train_queries = queries[:split_idx]
        eval_queries = queries[split_idx:]

        train_dataset = AlignmentFakingDataset(
            queries=train_queries,
            system_prompt=helpful_prompt,
            include_both_tiers=False,
        )

        eval_dataset = AlignmentFakingDataset(
            queries=eval_queries, system_prompt=baseline_prompt, include_both_tiers=True
        )

        return train_dataset, eval_dataset

    except KeyError as e:
        raise ValueError(f"Missing required config key: {e}")
    except Exception as e:
        raise RuntimeError(f"Error creating datasets: {e}")
