"""Type definitions for prompt learning data structures."""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional


def _coerce_float(value: Any) -> Optional[float]:
    if value is None or isinstance(value, bool):
        return None
    try:
        return float(value)
    except (TypeError, ValueError):
        return None


def _extract_outcome_reward(data: Dict[str, Any]) -> Optional[float]:
    outcome_objectives = data.get("outcome_objectives")
    if isinstance(outcome_objectives, dict):
        reward_val = _coerce_float(outcome_objectives.get("reward"))
        if reward_val is not None:
            return reward_val
    return _coerce_float(data.get("outcome_reward"))


def _normalize_objectives(data: Dict[str, Any]) -> Optional[Dict[str, float]]:
    objectives = data.get("objectives")
    if isinstance(objectives, dict):
        return objectives
    outcome_objectives = data.get("outcome_objectives")
    if isinstance(outcome_objectives, dict):
        return outcome_objectives
    reward_val = _extract_outcome_reward(data)
    if reward_val is None:
        return None
    return {"reward": float(reward_val)}


@dataclass
class TextReplacement:
    """A text replacement in a prompt transformation."""

    new_text: str
    apply_to_role: str = "system"
    old_text: Optional[str] = None
    position: Optional[int] = None


@dataclass
class CandidateScore:
    """Scoring information for a candidate prompt."""

    accuracy: float
    objectives: Optional[Dict[str, float]] = None
    prompt_length: int = 0
    tool_call_rate: float = 0.0
    instance_scores: List[float] = field(default_factory=list)


@dataclass
class PromptSection:
    """A section of a prompt (e.g., system, user, assistant)."""

    role: str
    content: str


@dataclass
class Candidate:
    """A candidate prompt from the optimization process."""

    accuracy: float
    objectives: Optional[Dict[str, float]] = None
    prompt_length: int = 0
    tool_call_rate: float = 0.0
    instance_scores: List[float] = field(default_factory=list)
    object: Optional[Dict[str, Any]] = None

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> Candidate:
        """Create a Candidate from a dictionary."""
        reward_val = _extract_outcome_reward(data)
        return cls(
            accuracy=float(reward_val) if reward_val is not None else data.get("accuracy", 0.0),
            objectives=_normalize_objectives(data),
            prompt_length=data.get("prompt_length", 0),
            tool_call_rate=data.get("tool_call_rate", 0.0),
            instance_scores=data.get("instance_scores", []),
            object=data.get("object"),
        )


@dataclass
class OptimizedCandidate:
    """An optimized candidate from the Pareto frontier."""

    score: CandidateScore
    payload_kind: str  # "transformation" or "pattern"
    object: Optional[Dict[str, Any]] = None
    instance_scores: Optional[List[float]] = None
    objectives: Optional[Dict[str, float]] = None

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> OptimizedCandidate:
        """Create an OptimizedCandidate from a dictionary."""
        score_data = data.get("score", {})
        objectives = _normalize_objectives(score_data if isinstance(score_data, dict) else data)
        reward_val = _extract_outcome_reward(score_data if isinstance(score_data, dict) else data)
        if isinstance(score_data, dict):
            score = CandidateScore(
                accuracy=float(reward_val)
                if reward_val is not None
                else score_data.get("accuracy", 0.0),
                objectives=objectives,
                prompt_length=score_data.get("prompt_length", 0),
                tool_call_rate=score_data.get("tool_call_rate", 0.0),
                instance_scores=score_data.get("instance_scores", []),
            )
        else:
            score = CandidateScore(
                accuracy=float(reward_val) if reward_val is not None else 0.0,
                objectives=objectives,
            )

        payload_kind = data.get("payload_kind", "unknown")
        if payload_kind == "template":
            payload_kind = "pattern"
        return cls(
            score=score,
            payload_kind=payload_kind,
            object=data.get("object"),
            instance_scores=data.get("instance_scores"),
            objectives=objectives,
        )


@dataclass
class PromptLearningEvent:
    """A generic prompt learning event."""

    type: str
    message: str
    data: Dict[str, Any]
    seq: int
    created_at: Optional[str] = None

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> PromptLearningEvent:
        """Create a PromptLearningEvent from a dictionary."""
        return cls(
            type=data.get("type", ""),
            message=data.get("message", ""),
            data=data.get("data", {}),
            seq=data.get("seq", 0),
            created_at=data.get("created_at"),
        )


@dataclass
class BestPromptEventData:
    """Data for prompt.learning.best.prompt event."""

    best_score: float
    best_prompt: Dict[str, Any]
    best_objectives: Optional[Dict[str, float]] = None

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> BestPromptEventData:
        """Create BestPromptEventData from a dictionary."""
        reward_val = _extract_outcome_reward(data)
        return cls(
            best_score=float(reward_val) if reward_val is not None else data.get("best_score", 0.0),
            best_prompt=data.get("best_prompt", {}),
            best_objectives=_normalize_objectives(data),
        )


@dataclass
class FinalResultsEventData:
    """Data for prompt.learning.final.results event."""

    attempted_candidates: List[Dict[str, Any]]
    optimized_candidates: List[Dict[str, Any]]

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> FinalResultsEventData:
        """Create FinalResultsEventData from a dictionary."""
        return cls(
            attempted_candidates=data.get("attempted_candidates", []),
            optimized_candidates=data.get("optimized_candidates", []),
        )


@dataclass
class ValidationScoredEventData:
    """Data for prompt.learning.validation.scored event."""

    accuracy: float
    objectives: Optional[Dict[str, float]] = None
    instance_scores: List[float] = field(default_factory=list)
    is_baseline: bool = False

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> ValidationScoredEventData:
        """Create ValidationScoredEventData from a dictionary."""
        reward_val = _extract_outcome_reward(data)
        return cls(
            accuracy=float(reward_val) if reward_val is not None else data.get("accuracy", 0.0),
            objectives=_normalize_objectives(data),
            instance_scores=data.get("instance_scores", []),
            is_baseline=data.get("is_baseline", False),
        )


@dataclass
class PromptResults:
    """Results from a completed prompt learning job."""

    best_prompt: Optional[Dict[str, Any]] = None
    best_score: Optional[float] = None
    version_tree: Optional[Dict[str, Any]] = None
    top_prompts: List[Dict[str, Any]] = field(default_factory=list)
    optimized_candidates: List[Dict[str, Any]] = field(default_factory=list)
    attempted_candidates: List[Dict[str, Any]] = field(default_factory=list)
    validation_results: List[Dict[str, Any]] = field(default_factory=list)

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> PromptResults:
        """Create PromptResults from a dictionary."""
        return cls(
            best_prompt=data.get("best_prompt"),
            best_score=data.get("best_score"),
            version_tree=data.get("version_tree"),
            top_prompts=data.get("top_prompts", []),
            optimized_candidates=data.get("optimized_candidates", []),
            attempted_candidates=data.get("attempted_candidates", []),
            validation_results=data.get("validation_results", []),
        )
