"""Benchmark tasks loader and manager for calibration."""

from __future__ import annotations

import logging
from dataclasses import dataclass
from pathlib import Path
from functools import lru_cache

import yaml

logger = logging.getLogger("ecate.calibration.tasks")

# Built-in tasks directory (relative to this module)
_BUILTIN_TASKS_DIR = Path(__file__).parent / "tasks"


@dataclass
class BenchmarkTask:
    """A single benchmark task for calibration."""
    id: str
    category: str
    difficulty: str  # simple | moderate | complex | expert
    prompt: str
    language: str  # en | it
    expected_behavior: str = ""
    deterministic_checks: list[dict] | None = None


def _load_yaml_file(path: Path) -> list[BenchmarkTask]:
    """Load tasks from a single YAML file."""
    tasks: list[BenchmarkTask] = []
    
    try:
        with open(path) as f:
            data = yaml.safe_load(f)
        
        if not isinstance(data, list):
            logger.warning(f"Expected list in {path}, got {type(data)}")
            return tasks
        
        for item in data:
            if not isinstance(item, dict):
                continue
            
            task_id = item.get("id", "")
            if not task_id:
                logger.warning(f"Task missing id in {path}")
                continue
            
            # Validate required fields
            required = ["category", "difficulty", "prompt", "language"]
            missing = [f for f in required if not item.get(f)]
            if missing:
                logger.warning(f"Task {task_id} missing fields: {missing}")
                continue
            
            # Validate difficulty
            if item["difficulty"] not in ("simple", "moderate", "complex", "expert"):
                logger.warning(f"Task {task_id} has invalid difficulty: {item['difficulty']}")
                continue
            
            # Validate language
            if item["language"] not in ("en", "it"):
                logger.warning(f"Task {task_id} has invalid language: {item['language']}")
                continue
            
            tasks.append(BenchmarkTask(
                id=task_id,
                category=item["category"],
                difficulty=item["difficulty"],
                prompt=item["prompt"],
                language=item["language"],
                expected_behavior=item.get("expected_behavior", ""),
                deterministic_checks=item.get("deterministic_checks"),
            ))
    
    except Exception as e:
        logger.error(f"Error loading {path}: {e}")
    
    return tasks


def load_tasks_from_directory(directory: str | Path) -> list[BenchmarkTask]:
    """Load benchmark tasks from all YAML files in a directory.

    Args:
        directory: Path to directory containing .yaml/.yml files

    Returns:
        List of BenchmarkTask objects
    """
    path = Path(directory)
    if not path.exists():
        logger.warning(f"Tasks directory not found: {path}")
        return []

    tasks: list[BenchmarkTask] = []
    seen_ids: set[str] = set()

    for yaml_file in sorted(path.glob("*.y*ml")):
        file_tasks = _load_yaml_file(yaml_file)
        
        for task in file_tasks:
            if task.id in seen_ids:
                logger.warning(f"Duplicate task id '{task.id}' in {yaml_file}, skipping")
                continue
            tasks.append(task)
            seen_ids.add(task.id)

    logger.debug(f"Loaded {len(tasks)} tasks from {path}")
    return tasks


@lru_cache(maxsize=1)
def get_builtin_tasks() -> tuple[BenchmarkTask, ...]:
    """Return all built-in benchmark tasks (cached).
    
    Returns a tuple for immutability. Use list() if you need to modify.
    """
    tasks = load_tasks_from_directory(_BUILTIN_TASKS_DIR)
    logger.info(f"Loaded {len(tasks)} built-in benchmark tasks")
    return tuple(tasks)


def load_custom_tasks(directory: str | Path) -> list[BenchmarkTask]:
    """Load custom benchmark tasks from a user-specified directory.

    This is an alias for load_tasks_from_directory with clearer naming.
    
    Expected YAML format:
    ```yaml
    - id: custom_001
      category: code
      difficulty: moderate
      language: en
      prompt: |
        Write a function that...
      expected_behavior: |
        Should return...
      deterministic_checks:  # optional
        - type: python_exec
          test_code: "assert func(1) == 1"
    ```
    """
    return load_tasks_from_directory(directory)


def get_all_tasks(custom_dir: str | Path | None = None) -> list[BenchmarkTask]:
    """Get built-in tasks plus optional custom tasks.
    
    Args:
        custom_dir: Optional path to custom tasks directory
        
    Returns:
        Combined list of all tasks
    """
    tasks = list(get_builtin_tasks())
    
    if custom_dir:
        custom = load_custom_tasks(custom_dir)
        # Check for ID conflicts
        builtin_ids = {t.id for t in tasks}
        for task in custom:
            if task.id in builtin_ids:
                logger.warning(f"Custom task '{task.id}' conflicts with built-in, skipping")
            else:
                tasks.append(task)
    
    return tasks


def get_tasks_by_category(category: str) -> list[BenchmarkTask]:
    """Filter built-in tasks by category."""
    return [t for t in get_builtin_tasks() if t.category == category]


def get_tasks_by_language(language: str) -> list[BenchmarkTask]:
    """Filter built-in tasks by language."""
    return [t for t in get_builtin_tasks() if t.language == language]


def get_tasks_by_difficulty(difficulty: str) -> list[BenchmarkTask]:
    """Filter built-in tasks by difficulty."""
    return [t for t in get_builtin_tasks() if t.difficulty == difficulty]


def get_task_stats() -> dict:
    """Return statistics about the built-in task suite."""
    tasks = get_builtin_tasks()
    stats = {
        "total": len(tasks),
        "by_category": {},
        "by_language": {},
        "by_difficulty": {},
    }
    for task in tasks:
        stats["by_category"][task.category] = stats["by_category"].get(task.category, 0) + 1
        stats["by_language"][task.language] = stats["by_language"].get(task.language, 0) + 1
        stats["by_difficulty"][task.difficulty] = stats["by_difficulty"].get(task.difficulty, 0) + 1
    return stats


def get_categories() -> list[str]:
    """Return list of all task categories."""
    return sorted(set(t.category for t in get_builtin_tasks()))


def get_task_by_id(task_id: str) -> BenchmarkTask | None:
    """Find a task by its ID."""
    for task in get_builtin_tasks():
        if task.id == task_id:
            return task
    return None
