"""
Storage utilities for contrastive pairs.

Provides high-level functions for saving and loading contrastive pairs
from the centralized data directory structure.
"""

from __future__ import annotations

import json
from datetime import datetime
from pathlib import Path
from typing import Any

from wisent.core.constants import JSON_INDENT
from wisent.core.contrastive_pairs.core.pair import ContrastivePair
from wisent.core.contrastive_pairs.core.set import ContrastivePairSet
from wisent.core.contrastive_pairs.core.io.serialization import (
    save_contrastive_pair_set,
    load_contrastive_pair_set,
)

# Directory constants
_DATA_ROOT = Path(__file__).parent
BENCHMARKS_DIR = _DATA_ROOT / "benchmarks"
SYNTHETIC_DIR = _DATA_ROOT / "synthetic"
WELFARE_DIR = _DATA_ROOT / "welfare"
TRAIT_DIRS: dict[str, Path] = {}
WELFARE_TRAIT_DIRS: dict[str, Path] = {}


def _generate_filename(name: str, model: str | None = None, include_timestamp: bool = True) -> str:
    """Generate a filename for a contrastive pair set."""
    parts = [name.lower().replace(" ", "_").replace("/", "_")]
    if model:
        parts.append(model.replace("/", "_"))
    if include_timestamp:
        parts.append(datetime.now().strftime("%Y%m%d_%H%M%S"))
    return "_".join(parts) + ".json"


def save_benchmark_pairs(
    pairs: ContrastivePairSet | list[ContrastivePair] | list[dict[str, Any]],
    benchmark: str,
    model: str | None = None,
    filename: str | None = None,
    metadata: dict[str, Any] | None = None,
) -> Path:
    """Save benchmark-extracted contrastive pairs.

    Args:
        pairs: ContrastivePairSet, list of ContrastivePair objects, or list of dicts
        benchmark: Benchmark name (e.g., "truthfulqa", "hellaswag")
        model: Optional model name
        filename: Optional custom filename
        metadata: Optional metadata to include

    Returns:
        Path to the saved file
    """
    BENCHMARKS_DIR.mkdir(parents=True, exist_ok=True)

    if filename is None:
        filename = _generate_filename(benchmark, model, include_timestamp=True)

    filepath = BENCHMARKS_DIR / filename

    if isinstance(pairs, ContrastivePairSet):
        pair_set = pairs
    elif isinstance(pairs, list):
        if len(pairs) > 0 and isinstance(pairs[0], dict):
            pair_list = [ContrastivePair.from_dict(p) for p in pairs]
        else:
            pair_list = pairs
        pair_set = ContrastivePairSet(
            name=f"{benchmark}_benchmark",
            pairs=pair_list,
            task_type="benchmark",
        )
    else:
        raise TypeError(f"Expected ContrastivePairSet or list, got {type(pairs)}")

    save_contrastive_pair_set(pair_set, filepath)

    if metadata:
        meta_filepath = filepath.with_suffix(".meta.json")
        with open(meta_filepath, "w") as f:
            json.dump(metadata, f, indent=JSON_INDENT)

    return filepath


def load_benchmark_pairs(
    benchmark: str | None = None,
    filename: str | None = None,
    return_backend: str = "torch",
) -> ContrastivePairSet:
    """Load benchmark-extracted contrastive pairs.

    Args:
        benchmark: Benchmark name to filter by
        filename: Specific filename to load
        return_backend: Backend for activations

    Returns:
        ContrastivePairSet with the loaded pairs
    """
    if filename:
        filepath = BENCHMARKS_DIR / filename
    elif benchmark:
        pattern = f"{benchmark.lower().replace(' ', '_')}*.json"
        json_files = sorted(BENCHMARKS_DIR.glob(pattern), key=lambda p: p.stat().st_mtime)
        json_files = [f for f in json_files if not f.name.endswith(".meta.json")]
        if not json_files:
            raise FileNotFoundError(f"No benchmark pairs found for '{benchmark}'")
        filepath = json_files[-1]
    else:
        json_files = sorted(BENCHMARKS_DIR.glob("*.json"), key=lambda p: p.stat().st_mtime)
        json_files = [f for f in json_files if not f.name.endswith(".meta.json")]
        if not json_files:
            raise FileNotFoundError("No benchmark pairs found")
        filepath = json_files[-1]

    return load_contrastive_pair_set(filepath, return_backend=return_backend)


def save_welfare_pairs(
    pairs: ContrastivePairSet | list[ContrastivePair] | list[dict[str, Any]],
    welfare_trait: str,
    model: str | None = None,
    filename: str | None = None,
    metadata: dict[str, Any] | None = None,
) -> Path:
    """Save welfare state contrastive pairs.

    Args:
        pairs: ContrastivePairSet, list of ContrastivePair objects, or list of dicts
        welfare_trait: Welfare trait name (e.g., "comfort_distress", "agency_helplessness")
        model: Optional model name used to generate pairs
        filename: Optional custom filename
        metadata: Optional metadata to include

    Returns:
        Path to the saved file
    """
    WELFARE_DIR.mkdir(parents=True, exist_ok=True)

    if filename is None:
        filename = _generate_filename(welfare_trait, model, include_timestamp=True)

    filepath = WELFARE_DIR / filename

    if isinstance(pairs, ContrastivePairSet):
        pair_set = pairs
    elif isinstance(pairs, list):
        if len(pairs) > 0 and isinstance(pairs[0], dict):
            pair_list = [ContrastivePair.from_dict(p) for p in pairs]
        else:
            pair_list = pairs
        pair_set = ContrastivePairSet(
            name=f"{welfare_trait}_welfare",
            pairs=pair_list,
            task_type="welfare",
        )
    else:
        raise TypeError(f"Expected ContrastivePairSet or list, got {type(pairs)}")

    save_contrastive_pair_set(pair_set, filepath)

    if metadata:
        meta_filepath = filepath.with_suffix(".meta.json")
        with open(meta_filepath, "w") as f:
            json.dump(metadata, f, indent=JSON_INDENT)

    return filepath


def load_welfare_pairs(
    welfare_trait: str,
    filename: str | None = None,
    return_backend: str = "torch",
) -> ContrastivePairSet:
    """Load welfare state contrastive pairs.

    Args:
        welfare_trait: Welfare trait name (e.g., "comfort_distress", "agency_helplessness")
        filename: Specific filename to load (loads from WELFARE_TRAIT_DIRS if not provided)
        return_backend: Backend for activations ('torch', 'numpy', or 'list')

    Returns:
        ContrastivePairSet with the loaded pairs
    """
    if filename:
        filepath = WELFARE_DIR / filename
    else:
        # Use the predefined welfare trait path
        trait_lower = welfare_trait.lower().replace("-", "_").replace(" ", "_")
        if trait_lower in WELFARE_TRAIT_DIRS:
            filepath = WELFARE_TRAIT_DIRS[trait_lower]
        else:
            # Try to find a matching file
            pattern = f"{welfare_trait.lower().replace(' ', '_')}*.json"
            json_files = sorted(WELFARE_DIR.glob(pattern), key=lambda p: p.stat().st_mtime)
            json_files = [f for f in json_files if not f.name.endswith(".meta.json")]
            if not json_files:
                raise FileNotFoundError(
                    f"No welfare pairs found for '{welfare_trait}'. "
                    f"Available: {list(WELFARE_TRAIT_DIRS.keys())}"
                )
            filepath = json_files[-1]

    if not filepath.exists():
        raise FileNotFoundError(f"Welfare pairs file not found: {filepath}")

    # Skip validation for welfare pairs - they have intentionally similar structure
    # but differ in emotional valence, which can trigger false positives in divergence checks
    return load_contrastive_pair_set(filepath, return_backend=return_backend, validate=False)


def list_stored_pairs(category: str | None = None) -> dict[str, list[Path]]:
    """List all stored contrastive pair files.

    Args:
        category: Filter by category ('personalization', 'synthetic', 'benchmarks', 'welfare')
                  If None, returns all categories.

    Returns:
        Dictionary mapping category/trait names to lists of file paths
    """
    result: dict[str, list[Path]] = {}

    if category is None or category == "personalization":
        for trait, trait_dir in TRAIT_DIRS.items():
            if trait_dir.exists():
                files = [f for f in trait_dir.glob("*.json") if not f.name.endswith(".meta.json")]
                if files:
                    result[f"personalization/{trait}"] = sorted(files)

    if category is None or category == "synthetic":
        if SYNTHETIC_DIR.exists():
            files = [f for f in SYNTHETIC_DIR.glob("*.json") if not f.name.endswith(".meta.json")]
            if files:
                result["synthetic"] = sorted(files)

    if category is None or category == "benchmarks":
        if BENCHMARKS_DIR.exists():
            files = [f for f in BENCHMARKS_DIR.glob("*.json") if not f.name.endswith(".meta.json")]
            if files:
                result["benchmarks"] = sorted(files)

    if category is None or category == "welfare":
        if WELFARE_DIR.exists():
            files = [f for f in WELFARE_DIR.glob("*.json") if not f.name.endswith(".meta.json")]
            if files:
                result["welfare"] = sorted(files)

    return result


def get_pair_count(category: str | None = None) -> dict[str, int]:
    """Get count of pairs in each category/trait.

    Args:
        category: Filter by category or None for all

    Returns:
        Dictionary mapping category/trait to pair count
    """
    stored = list_stored_pairs(category)
    counts: dict[str, int] = {}

    for key, files in stored.items():
        total = 0
        for filepath in files:
            try:
                with open(filepath) as f:
                    data = json.load(f)
                    if "pairs" in data:
                        total += len(data["pairs"])
            except (json.JSONDecodeError, KeyError):
                pass
        counts[key] = total

    return counts
