"""Dataset loader for JSONL and CSV formats."""

from __future__ import annotations

import csv
import json
import random
from pathlib import Path
from typing import Optional, Any

from .registry import get_builtin_path, BUILTIN_DATASETS
from ..utils import get_logger

logger = get_logger(__name__)


class DatasetLoader:
    """
    Loader for evaluation datasets.
    
    Supports:
    - Built-in datasets (by name)
    - JSONL files
    - CSV files
    
    Canonical schema:
    - id: str (required)
    - prompt: str (required)  
    - dataset_name: str (auto-filled from filename if missing)
    - category: Optional[str]
    - reference: Optional[str]
    - competitor_output: Optional[str]
    - metadata: Optional[dict]
    """
    
    def load(
        self,
        source: str,
        max_samples: Optional[int] = None,
        seed: Optional[int] = None,
    ) -> list[dict[str, Any]]:
        """
        Load prompts from a source.
        
        Args:
            source: Builtin name or path to JSONL/CSV file
            max_samples: Maximum number of samples to load
            seed: Random seed for sampling (None for deterministic order)
            
        Returns:
            List of prompt dicts with canonical schema
        """
        # Check if builtin
        if source in BUILTIN_DATASETS:
            path = get_builtin_path(source)
        else:
            path = Path(source)
            if not path.exists():
                raise FileNotFoundError(f"Dataset not found: {source}")
        
        # Load based on extension
        ext = path.suffix.lower()
        if ext == ".jsonl":
            prompts = self._load_jsonl(path)
        elif ext == ".csv":
            prompts = self._load_csv(path)
        else:
            raise ValueError(f"Unsupported file format: {ext}. Use .jsonl or .csv")
        
        # Auto-fill dataset_name if missing
        dataset_name = path.stem
        for p in prompts:
            if not p.get("dataset_name"):
                p["dataset_name"] = dataset_name
        
        # Sample if needed
        if max_samples and len(prompts) > max_samples:
            if seed is not None:
                random.seed(seed)
                prompts = random.sample(prompts, max_samples)
            else:
                prompts = prompts[:max_samples]
        
        logger.info(f"Loaded {len(prompts)} prompts from {source}")
        return prompts
    
    def load_category(
        self,
        category: str,
        max_samples: Optional[int] = None,
        seed: Optional[int] = None,
    ) -> list[dict[str, Any]]:
        """
        Load prompts for an evaluation category, aggregating from multiple sources.
        
        Args:
            category: Category name (e.g., "physics", "chemistry", "law")
            max_samples: Maximum total samples to load
            seed: Random seed for sampling
            
        Returns:
            List of prompt dicts from all category sources
        """
        from .categories import get_category_sources
        
        sources = get_category_sources(category)
        if not sources:
            raise ValueError(f"No source files found for category: {category}")
        
        all_prompts = []
        for source_path in sources:
            try:
                prompts = self._load_jsonl(source_path)
                # Tag prompts with source info
                for p in prompts:
                    if not p.get("dataset_name"):
                        p["dataset_name"] = source_path.stem
                    p["source_category"] = category
                all_prompts.extend(prompts)
            except Exception as e:
                logger.warning(f"Failed to load {source_path}: {e}")
        
        # Sample if needed
        if max_samples and len(all_prompts) > max_samples:
            import random
            if seed is not None:
                random.seed(seed)
                all_prompts = random.sample(all_prompts, max_samples)
            else:
                all_prompts = all_prompts[:max_samples]
        
        logger.info(f"Loaded {len(all_prompts)} prompts for category '{category}' from {len(sources)} sources")
        return all_prompts
    
    def _load_jsonl(self, path: Path) -> list[dict[str, Any]]:
        """Load JSONL file."""
        prompts = []
        with open(path, "r", encoding="utf-8") as f:
            for i, line in enumerate(f):
                line = line.strip()
                if not line:
                    continue
                try:
                    data = json.loads(line)
                    prompt = self._normalize(data, i)
                    prompts.append(prompt)
                except json.JSONDecodeError as e:
                    logger.warning(f"Invalid JSON at line {i+1}: {e}")
        return prompts
    
    def _load_csv(self, path: Path) -> list[dict[str, Any]]:
        """Load CSV file."""
        prompts = []
        with open(path, "r", encoding="utf-8", newline="") as f:
            reader = csv.DictReader(f)
            for i, row in enumerate(reader):
                prompt = self._normalize(dict(row), i)
                prompts.append(prompt)
        return prompts
    
    def _normalize(self, data: dict[str, Any], index: int) -> dict[str, Any]:
        """Normalize to canonical schema."""
        # Handle different column names for prompt
        prompt_text = (
            data.get("prompt") or 
            data.get("text") or 
            data.get("input") or
            data.get("question") or
            ""
        )
        
        # Handle ID
        prompt_id = str(
            data.get("id") or 
            data.get("prompt_id") or 
            f"prompt_{index}"
        )
        
        return {
            "id": prompt_id,
            "prompt": prompt_text,
            "dataset_name": data.get("dataset_name"),
            "category": data.get("category"),
            "reference": data.get("reference") or data.get("answer") or data.get("expected"),
            "competitor_output": data.get("competitor_output") or data.get("baseline"),
            "metadata": data.get("metadata"),
        }
