"""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 _coerce_int(value: Any) -> Optional[int]:
    if value is None or isinstance(value, bool):
        return None
    try:
        return int(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)}


def _infer_lever_versions_from_summary(lever_summary: Optional[Dict[str, Any]]) -> Dict[str, int]:
    if not isinstance(lever_summary, dict):
        return {}
    prompt_lever_id = lever_summary.get("prompt_lever_id")
    candidate_versions = lever_summary.get("candidate_lever_versions")
    if not isinstance(prompt_lever_id, str) or not isinstance(candidate_versions, dict):
        return {}
    candidate_id = (
        lever_summary.get("selected_candidate_id")
        or lever_summary.get("best_candidate_id")
        or (next(iter(candidate_versions.keys())) if len(candidate_versions) == 1 else None)
    )
    if candidate_id is None:
        return {}
    version = _coerce_int(candidate_versions.get(str(candidate_id)))
    if version is None:
        return {}
    return {prompt_lever_id: version}


def _infer_best_lever_version_from_summary(
    lever_summary: Optional[Dict[str, Any]],
    *,
    best_candidate_id: Optional[str] = None,
) -> Optional[int]:
    if not isinstance(lever_summary, dict):
        return None
    candidate_versions = lever_summary.get("candidate_lever_versions")
    if not isinstance(candidate_versions, dict):
        return None
    candidate_id = (
        best_candidate_id
        or lever_summary.get("selected_candidate_id")
        or lever_summary.get("best_candidate_id")
        or (next(iter(candidate_versions.keys())) if len(candidate_versions) == 1 else None)
    )
    if candidate_id is None:
        return None
    return _coerce_int(candidate_versions.get(str(candidate_id)))


@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 CandidateReward:
    """Reward information for a candidate prompt."""

    accuracy: float
    objectives: Optional[Dict[str, float]] = None
    prompt_length: int = 0
    tool_call_rate: float = 0.0
    instance_rewards: 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_rewards: List[float] = field(default_factory=list)
    object: Optional[Dict[str, Any]] = None
    seed_rewards: Optional[List[Dict[str, Any]]] = None  # Per-seed rewards: [{seed, reward}, ...]

    @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_rewards=data.get("instance_rewards", []),
            object=data.get("object"),
            seed_rewards=data.get("seed_rewards"),
        )


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

    reward: CandidateReward
    payload_kind: str  # "transformation" or "pattern"
    object: Optional[Dict[str, Any]] = None
    instance_rewards: 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."""
        reward_data = data.get("reward", {})
        objectives = _normalize_objectives(reward_data if isinstance(reward_data, dict) else data)
        reward_val = _extract_outcome_reward(reward_data if isinstance(reward_data, dict) else data)
        if isinstance(reward_data, dict):
            reward = CandidateReward(
                accuracy=float(reward_val)
                if reward_val is not None
                else reward_data.get("accuracy", 0.0),
                objectives=objectives,
                prompt_length=reward_data.get("prompt_length", 0),
                tool_call_rate=reward_data.get("tool_call_rate", 0.0),
                instance_rewards=reward_data.get("instance_rewards", []),
            )
        else:
            reward = CandidateReward(
                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(
            reward=reward,
            payload_kind=payload_kind,
            object=data.get("object"),
            instance_rewards=data.get("instance_rewards"),
            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_reward: float
    best_candidate: 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_reward=float(reward_val)
            if reward_val is not None
            else data.get("best_reward", 0.0),
            best_candidate=data.get("best_candidate") or data.get("best_prompt") or {},
            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_rewards: 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_rewards=data.get("instance_rewards", []),
            is_baseline=data.get("is_baseline", False),
        )


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

    best_candidate: Optional[Dict[str, Any] | str] = None  # Best candidate payload (when available)
    best_candidate_content: Optional[str] = None
    best_reward: Optional[float] = None
    version_tree: Optional[Dict[str, Any]] = None
    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)
    total_rollouts: int = 0  # Total number of rollouts (metric evaluations)
    total_proposal_calls: int = 0  # Total number of proposal/mutation calls (LLM reflection calls)
    lever_summary: Optional[Dict[str, Any]] = None
    sensor_frames: List[Dict[str, Any]] = field(default_factory=list)
    lever_versions: Dict[str, int] = field(default_factory=dict)
    best_lever_version: Optional[int] = None
    event_counts: Dict[str, int] = field(default_factory=dict)
    event_history: List[Dict[str, Any]] = field(default_factory=list)
    gepa: Dict[str, Any] = field(default_factory=dict)
    mipro: Dict[str, Any] = field(default_factory=dict)

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> PromptResults:
        """Create PromptResults from a dictionary."""
        lever_summary = (
            data.get("lever_summary") if isinstance(data.get("lever_summary"), dict) else None
        )
        lever_versions: Dict[str, int] = {}
        raw_versions = data.get("lever_versions")
        if isinstance(raw_versions, dict):
            for key, value in raw_versions.items():
                try:
                    lever_versions[str(key)] = int(value)
                except (TypeError, ValueError):
                    continue
        inferred_versions = _infer_lever_versions_from_summary(lever_summary)
        summary_prompt_lever_id = (
            lever_summary.get("prompt_lever_id") if isinstance(lever_summary, dict) else None
        )
        if inferred_versions and (
            not lever_versions
            or (
                isinstance(summary_prompt_lever_id, str)
                and summary_prompt_lever_id not in lever_versions
            )
        ):
            lever_versions = inferred_versions
        best_lever_version_raw = data.get("best_lever_version")
        best_lever_version: Optional[int] = None
        if best_lever_version_raw is not None:
            try:
                best_lever_version = int(best_lever_version_raw)
            except (TypeError, ValueError):
                best_lever_version = None
        if best_lever_version is None:
            best_lever_version = _infer_best_lever_version_from_summary(lever_summary)
        if best_lever_version is None and lever_versions:
            best_lever_version = max(lever_versions.values())
        return cls(
            best_candidate=data.get("best_candidate") or data.get("best_prompt"),
            best_candidate_content=data.get("best_candidate_content"),
            best_reward=data.get("best_reward"),
            version_tree=data.get("version_tree"),
            optimized_candidates=data.get("optimized_candidates", []),
            attempted_candidates=data.get("attempted_candidates", []),
            validation_results=data.get("validation_results", []),
            total_rollouts=data.get("total_rollouts", 0),
            total_proposal_calls=data.get("total_proposal_calls", 0),
            lever_summary=lever_summary,
            sensor_frames=[
                frame for frame in data.get("sensor_frames", []) if isinstance(frame, dict)
            ]
            if isinstance(data.get("sensor_frames"), list)
            else [],
            lever_versions=lever_versions,
            best_lever_version=best_lever_version,
            event_counts=data.get("event_counts", {}),
            event_history=data.get("event_history", []),
            gepa=data.get("gepa", {}),
            mipro=data.get("mipro", {}),
        )
