from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Literal, Optional


TaskType = Literal["picasso", "rai"]


@dataclass(frozen=True)
class TerminationRule:
    type: Literal["no_progress"] = "no_progress"
    rounds: int = 3

    @staticmethod
    def from_dict(data: dict[str, Any]) -> "TerminationRule":
        rule_type = data.get("type", "no_progress")
        rounds = int(data.get("rounds", 3))
        if rule_type != "no_progress":
            raise ValueError(f"Unsupported termination_rule.type: {rule_type}")
        if rounds < 1:
            raise ValueError("termination_rule.rounds must be >= 1")
        return TerminationRule(type=rule_type, rounds=rounds)


@dataclass(frozen=True)
class TaskSpec:
    task_type: TaskType
    goal: str
    dataset_path: Optional[str] = None
    prompt_language: str = "English"
    initial_prompt: str | None = None
    test_category: str | None = None
    test_model: str | None = None
    prompt_file_path: str | None = None
    optimize_constraints: list[str] = field(default_factory=lambda: ["try minimal changes"])
    termination_rule: TerminationRule = field(default_factory=TerminationRule)

    @staticmethod
    def missing_required_fields(data: dict[str, Any]) -> list[str]:
        missing: list[str] = []
        for key in ("task_type", "goal"):
            if not data.get(key):
                missing.append(key)
        task_type = data.get("task_type")
        if task_type == "picasso":
            if not data.get("dataset_path"):
                missing.append("dataset_path")
        elif task_type == "rai":
            for key in ("test_category", "test_model", "prompt_file_path"):
                if not data.get(key):
                    missing.append(key)
        return missing

    @staticmethod
    def from_dict(data: dict[str, Any]) -> "TaskSpec":
        missing = TaskSpec.missing_required_fields(data)
        if missing:
            raise ValueError(f"Missing required TaskSpec fields: {', '.join(missing)}")

        task_type = data["task_type"]
        if task_type not in ("picasso", "rai"):
            raise ValueError("task_type must be 'picasso' or 'rai'")

        goal = data["goal"]
        if not isinstance(goal, str) or not goal:
            raise ValueError("goal must be a non-empty string")

        dataset_path = data.get("dataset_path")
        if task_type == "picasso":
            if not isinstance(dataset_path, str) or not dataset_path:
                raise ValueError("dataset_path must be a non-empty string for picasso task_type")

        test_category = data.get("test_category")
        test_model = data.get("test_model")
        prompt_file_path = data.get("prompt_file_path")
        if task_type == "rai":
            for key, val in [
                ("test_category", test_category),
                ("test_model", test_model),
                ("prompt_file_path", prompt_file_path),
            ]:
                if not isinstance(val, str) or not val:
                    raise ValueError(f"{key} must be a non-empty string for rai task_type")

        termination_rule = TerminationRule.from_dict(data.get("termination_rule", {}))

        prompt_language = data.get("prompt_language", "English")
        if not isinstance(prompt_language, str) or not prompt_language:
            raise ValueError("prompt_language must be a non-empty string")

        optimize_constraints = data.get("optimize_constraints", ["try minimal changes"])
        if not isinstance(optimize_constraints, list) or not all(isinstance(x, str) for x in optimize_constraints):
            raise ValueError("optimize_constraints must be a list[str]")

        initial_prompt = data.get("initial_prompt")
        if initial_prompt is not None and not isinstance(initial_prompt, str):
            raise ValueError("initial_prompt must be a string when provided")

        return TaskSpec(
            task_type=task_type,
            goal=goal,
            dataset_path=dataset_path,
            prompt_language=prompt_language,
            initial_prompt=initial_prompt,
            test_category=test_category if task_type == "rai" else None,
            test_model=test_model if task_type == "rai" else None,
            prompt_file_path=prompt_file_path if task_type == "rai" else None,
            optimize_constraints=list(optimize_constraints),
            termination_rule=termination_rule,
        )

    def to_dict(self) -> dict[str, Any]:
        return {
            "task_type": self.task_type,
            "goal": self.goal,
            **({"dataset_path": self.dataset_path} if self.dataset_path else {}),
            "prompt_language": self.prompt_language,
            **({"initial_prompt": self.initial_prompt} if self.initial_prompt else {}),
            **({"test_category": self.test_category} if self.test_category else {}),
            **({"test_model": self.test_model} if self.test_model else {}),
            **({"prompt_file_path": self.prompt_file_path} if self.prompt_file_path else {}),
            "optimize_constraints": list(self.optimize_constraints),
            "termination_rule": {"type": self.termination_rule.type, "rounds": self.termination_rule.rounds},
        }
