import json
from pathlib import Path
from typing import Any, Callable, Dict, Optional, Tuple, TypedDict

from pydantic import BaseModel, Field

from dojo_sdk_core.dojos.rewards import REWARD_FUNCTIONS
from dojo_sdk_core.dojos.rewards.backend import Backend


class RewardNotFoundError(Exception):
    """Exception raised when a reward function is not found."""

    pass


# V2 reward types
class StateKeyQuery(TypedDict):
    collection: str
    filter: Dict[str, Any]


StateKey = Dict[str, StateKeyQuery]
ValidatorFunc = Callable[[Dict[str, Any], Dict[str, Any]], Tuple[float, str]]


class ValidateTask(TypedDict):
    state_key: StateKey
    validate_backend: ValidatorFunc
    validate_frontend: ValidatorFunc


class SettingsConfig(BaseModel):
    """Settings configuration."""

    anthropic_api_key: str = Field("", description="Anthropic API key")
    openai_api_key: str = Field("", description="OpenAI API key")
    openai_api_url: str = Field("", description="OpenAI API URL")
    dojo_websocket_endpoint: str = Field("", description="Dojo websocket endpoint")
    dojo_http_endpoint: str = Field("", description="Dojo http endpoint")
    posthog_api_key: str = Field("", description="PostHog API key")


class EnvironmentConfig(BaseModel):
    """Environment configuration."""

    type: str = Field(..., description="Environment type (e.g., 'spa')")
    path: str = Field(..., description="Path to environment file")


class InstructionsConfig(BaseModel):
    """Task instructions configuration."""

    user_prompt: str = Field(..., description="Prompt to show to the agent")
    success_criteria: str = Field(..., description="What constitutes success")


class TaskDefinition(BaseModel):
    spa: str = Field(..., description="SPA name")
    id: str = Field(..., description="Unique task identifier")
    name: str = Field(..., description="Human-readable task name")
    description: str = Field(..., description="Task description")
    version: str = Field(..., description="Task version")
    env_version: str = Field(..., description="Environment version")
    initial_backend_state_name: str = Field(..., description="Initial backend state/dataset name")
    environment: EnvironmentConfig = Field(..., description="Environment configuration")
    initial_state: dict[str, Any] = Field(..., description="Initial state for the environment")
    instructions: InstructionsConfig = Field(..., description="Task instructions")
    reward: Optional[ValidateTask] = Field(
        None,
        description="ValidateTask bundle with state_key, validate_backend, validate_frontend",
    )
    reward_v2: Optional[Callable[[Dict[str, Any], Backend], Tuple[float, str]]] = Field(
        None,
        description="Reward function for v2 tasks",
    )
    reward_function_name: Optional[str] = Field(
        None,
        description="Name of the reward function (e.g., 'notion/task_create_page') for verifier service",
    )
    max_steps: int = Field(default=10, description="Maximum number of steps allowed")
    timeout_seconds: int = Field(default=60, description="Task timeout in seconds")
    metadata: dict[str, Any] = Field(default_factory=dict, description="Additional metadata")
    environment_type: str = Field(default="gui", description="Environment type e.g 'mcp' or 'gui'")
    trace_id: Optional[str] = Field(None, description="Unique trace identifier for the task")

    def model_post_init(self, __context: Any) -> None:
        """Validate that reward is provided."""
        if self.version == "2.0" and not self.reward_v2:
            raise ValueError("reward_v2 must be provided for v2 tasks")
        if self.version == "1.0" and not self.reward:
            raise ValueError("reward must be provided for v1 tasks")

    def get_environment_path(self, base_path: Optional[Path] = None) -> Path:
        """Get absolute path to environment file."""
        if base_path is None:
            base_path = Path.cwd()
        return base_path / self.environment.path

    @classmethod
    def from_hf_row(
        cls, row: dict[str, Any], reward_importer: Optional[Callable[[str], Optional[ValidateTask]]] = None
    ) -> "TaskDefinition":
        """Create a TaskDefinition from a HuggingFace dataset row."""
        initial_state = json.loads(row["initial_state"])
        environment = json.loads(row["environment"])
        instructions = json.loads(row["instructions"])
        metadata = json.loads(row["metadata"])
        version = row["version"]
        reward_v2 = None
        reward = None
        environment_type = row.get("environment_type")
        if not environment_type:
            print(f"environment_type is empty for task {row['id']}, setting to default 'gui'")
            environment_type = "gui"

        reward_function_name = row.get("reward_function", "").strip() if row.get("reward_function") else None

        if version == "2.0":
            reward_v2 = REWARD_FUNCTIONS.get(row["reward_function"].strip())
            if not reward_v2:
                raise RewardNotFoundError(f"Reward function {row['reward_function']} not found")
        elif version == "1.0" and row["reward_function"] and row["reward_function"].strip():
            if reward_importer:
                reward = reward_importer(row["reward_function"])

        return cls(
            id=row["id"],
            spa=row["spa"],
            name=row["name"],
            description=row["description"],
            environment=environment,
            initial_state=initial_state,
            version=version,
            env_version=row.get("env_version", ""),
            initial_backend_state_name=row.get("initial_backend_state_name", ""),
            instructions=instructions,
            reward=reward,
            reward_v2=reward_v2,
            reward_function_name=reward_function_name,
            max_steps=row["max_steps"],
            timeout_seconds=row["timeout_seconds"],
            metadata=metadata,
            environment_type=environment_type,
            trace_id=row.get("trace_id"),
        )
