from __future__ import annotations

import json
from dataclasses import dataclass
from pathlib import Path

from data_agent_baseline.benchmark.schema import PublicTask, TaskAssets, TaskRecord

TASK_DIR_PREFIX = "task_"


def _task_number(task_id: str) -> int:
    try:
        if not task_id.startswith(TASK_DIR_PREFIX):
            # raise ValueError(f"Invalid task id: {task_id}")
            return float("inf")
        return int(task_id.removeprefix(TASK_DIR_PREFIX))
    except ValueError:
        return float("inf")


def _load_task_record(task_json_path: Path) -> TaskRecord:
    payload = json.loads(task_json_path.read_text())
    required_keys = {"task_id", "question"}
    actual_keys = set(payload)
    missing_required = required_keys - actual_keys
    if missing_required:
        raise ValueError(
            f"Missing required task.json keys for {task_json_path.parent.name}: "
            f"{sorted(missing_required)}"
        )

    task_id = str(payload.get("task_id", "")).strip()
    question = str(payload.get("question", "")).strip()
    raw_difficulty = payload.get("difficulty", "unknown")
    difficulty = str(raw_difficulty).strip() or "unknown"

    if not task_id:
        raise ValueError(f"Invalid task_id in {task_json_path}: expected non-empty string")
    if not question:
        raise ValueError(f"Invalid question in {task_json_path}: expected non-empty string")

    return TaskRecord(
        task_id=task_id,
        difficulty=difficulty,
        question=question,
    )


@dataclass(frozen=True, slots=True)
class DABenchPublicDataset:
    root_dir: Path

    @property
    def exists(self) -> bool:
        return self.root_dir.is_dir()

    def task_dirs(self) -> list[Path]:
        if not self.exists:
            return []

        task_dirs = [
            path
            for path in self.root_dir.iterdir()
            if path.is_dir() and path.name.startswith(TASK_DIR_PREFIX)
        ]
        task_dirs.sort(key=lambda path: _task_number(path.name))
        return task_dirs

    def list_task_ids(self) -> list[str]:
        return [path.name for path in self.task_dirs()]

    def get_task(self, task_id: str) -> PublicTask:
        task_dir = self.root_dir / task_id
        task_json_path = task_dir / "task.json"
        if not task_json_path.exists():
            raise FileNotFoundError(f"Missing task.json: {task_json_path}")

        record = _load_task_record(task_json_path)
        if record.task_id != task_dir.name:
            raise ValueError(f"task_id mismatch for {task_dir}: task.json has {record.task_id}")

        context_dir = task_dir / "context"
        if not context_dir.is_dir():
            raise FileNotFoundError(f"Missing context dir: {context_dir}")

        assets = TaskAssets(task_dir=task_dir, context_dir=context_dir)
        return PublicTask(record=record, assets=assets)

    def iter_tasks(
        self,
        *,
        task_ids: list[str] | None = None,
        difficulty: str | None = None,
        difficulties: list[str] | None = None,
    ) -> list[PublicTask]:
        selected_ids = set(task_ids or [])
        selected_difficulties = set(difficulties or [])
        if difficulty is not None:
            selected_difficulties.add(difficulty)

        tasks: list[PublicTask] = []
        for task_dir in self.task_dirs():
            if selected_ids and task_dir.name not in selected_ids:
                continue
            task = self.get_task(task_dir.name)
            if selected_difficulties and task.difficulty not in selected_difficulties:
                continue
            tasks.append(task)
        return tasks

    def task_counts(self) -> dict[str, int]:
        counts: dict[str, int] = {}
        for task in self.iter_tasks():
            counts[task.difficulty] = counts.get(task.difficulty, 0) + 1
        return counts