"""Data models for improvement mechanisms."""
from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any

from .utils import utcnow


IMPROVEMENT_SCHEMA_KIND = "improvement_memory"
IMPROVEMENT_SCHEMA_VERSION = 1


@dataclass
class RewardComponent:
    """Single factor inside a unified reward score."""

    name: str
    raw_value: float
    weight: float
    contribution: float
    metadata: dict[str, Any] = field(default_factory=dict)


@dataclass
class RewardScore:
    """Unified reward score for a run or example."""

    score: float
    trace_id: str | None = None
    run_id: str | None = None
    workspace_id: str | None = None
    agent_name: str | None = None
    components: list[RewardComponent] = field(default_factory=list)
    created_at: str = field(default_factory=utcnow)


@dataclass
class LearnedPreferenceModelState:
    """Persistable learned preference scorer parameters."""

    version: int
    bias: float
    feature_weights: dict[str, float]
    training_summary: dict[str, Any] = field(default_factory=dict)


@dataclass
class PolicyDecision:
    """Learned policy recommendation for routing, retrieval, and tool choice."""

    workspace_id: str | None
    task_type: str | None
    user_segment: str | None
    model_key: str | None = None
    retrieval_limit: int | None = None
    retrieval_filter: dict[str, Any] | None = None
    preferred_tools: list[str] = field(default_factory=list)
    fallback_provider_keys: list[str] = field(default_factory=list)
    confidence: float = 0.0
    average_reward: float = 0.0
    rationale: list[str] = field(default_factory=list)
    metrics: dict[str, Any] = field(default_factory=dict)


@dataclass
class TrainingExport:
    """Exported training data for offline analysis."""

    workspace_id: str | None
    agent_name: str | None
    examples: list[dict[str, Any]]
    generated_at: str = field(default_factory=utcnow)
    metadata: dict[str, Any] = field(default_factory=dict)


@dataclass
class CriticIssue:
    """Structured defect identified by a critic."""

    kind: str
    severity: str
    description: str


@dataclass
class CriticReview:
    """Critic output over a candidate answer."""

    should_revise: bool
    summary: str
    issues: list[CriticIssue] = field(default_factory=list)
    prompt_patch: str | None = None


@dataclass
class RevisionSuggestion:
    """Improved answer or prompt proposed by a reviser."""

    improved_answer: str | None = None
    prompt_patch: str | None = None
    rationale: str | None = None


@dataclass
class CriticReviserResult:
    """Full critic-reviser cycle result."""

    review: CriticReview
    revision: RevisionSuggestion
    revised: bool


@dataclass
class ReflectionLesson:
    """Structured lesson learned from runs or mined patterns."""

    id: str
    workspace_id: str | None
    agent_name: str | None
    category: str
    lesson: str
    rationale: str
    suggested_action: str
    confidence: float
    source_run_ids: list[str] = field(default_factory=list)
    metadata: dict[str, Any] = field(default_factory=dict)
    created_at: str = field(default_factory=utcnow)


@dataclass
class FailurePattern:
    """Recurring failure pattern identified from runs."""

    id: str
    pattern_type: str
    description: str
    failure_types: list[str]
    occurrence_count: int
    affected_runs: list[str]
    suggested_fix: str | None = None
    suggested_fix_type: str = ""  # For test compatibility
    confidence: float = 0.0
    signature: str = ""  # For backward compatibility
    created_at: str = field(default_factory=lambda: __import__('datetime').datetime.now(__import__('datetime').UTC).isoformat())


@dataclass
class RootCauseHypothesis:
    """Hypothesis about root cause of failures."""

    id: str
    description: str
    evidence: list[str]
    confidence: float
    affected_runs: list[str]
    suggested_actions: list[str] = field(default_factory=list)
    created_at: str = field(default_factory=utcnow)


@dataclass
class RootCauseReport:
    """Aggregated root cause analysis report."""

    workspace_id: str | None
    agent_name: str | None
    window_days: int = 7
    total_failures: int = 0
    sample_size: int = 0
    automatic_classification_count: int = 0
    hypotheses: list[RootCauseHypothesis] = field(default_factory=list)
    top_failure_types: list[dict[str, Any]] = field(default_factory=list)
    failure_summary: dict[str, int] = field(default_factory=dict)
    correlations: dict[str, list[dict[str, Any]]] = field(default_factory=dict)
    recommendations: list[str] = field(default_factory=list)
    generated_at: str = field(default_factory=lambda: __import__('datetime').datetime.now(__import__('datetime').UTC).isoformat())


__all__ = [
    "IMPROVEMENT_SCHEMA_KIND",
    "IMPROVEMENT_SCHEMA_VERSION",
    "RewardComponent",
    "RewardScore",
    "LearnedPreferenceModelState",
    "PolicyDecision",
    "TrainingExport",
    "CriticIssue",
    "CriticReview",
    "RevisionSuggestion",
    "CriticReviserResult",
    "ReflectionLesson",
    "FailurePattern",
    "RootCauseHypothesis",
    "RootCauseReport",
]
