"""Pydantic models used throughout the ACET toolkit."""

from __future__ import annotations

import uuid
from datetime import datetime, timezone
from enum import Enum
from typing import List

from pydantic import BaseModel, ConfigDict, Field


class DeltaStatus(str, Enum):
    """Lifecycle status for a context delta."""

    STAGED = "staged"
    ACTIVE = "active"
    ARCHIVED = "archived"


class ContextDelta(BaseModel):
    """Atomic unit of learned context captured by ACET."""

    id: str = Field(
        default_factory=lambda: str(uuid.uuid4()),
        description="Unique identifier for the delta.",
    )
    topic: str = Field(..., description="Category or topic for this delta.")
    guideline: str = Field(..., description="The actionable rule or insight.")
    conditions: List[str] = Field(
        default_factory=list,
        description="Conditions under which this guideline applies.",
    )
    evidence: List[str] = Field(
        default_factory=list,
        description="Evidence or supporting references.",
    )
    tags: List[str] = Field(
        default_factory=list,
        description="Free-form tags for additional classification.",
    )

    version: int = Field(default=1, ge=1, description="Version of the delta.")
    status: DeltaStatus = Field(
        default=DeltaStatus.STAGED,
        description="Lifecycle status of the delta.",
    )

    score: float = Field(
        default=0.0,
        ge=0.0,
        le=1.0,
        description="Composite ranking score assigned by the curator.",
    )
    recency: float = Field(
        default=1.0,
        ge=0.0,
        le=1.0,
        description="Recency weighting factor.",
    )
    usage_count: int = Field(
        default=0,
        ge=0,
        description="Number of times this delta has been injected.",
    )
    helpful_count: int = Field(
        default=0,
        ge=0,
        description="Count of times the delta was marked helpful.",
    )
    harmful_count: int = Field(
        default=0,
        ge=0,
        description="Count of times the delta was marked harmful.",
    )

    author: str = Field(
        default="reflector",
        description="Origin of the delta (reflector, human reviewer, etc.).",
    )
    risk_level: str = Field(
        default="low",
        description="Risk classification (low, medium, high).",
    )
    confidence: float = Field(
        default=0.0,
        ge=0.0,
        le=1.0,
        description="Confidence score associated with the delta.",
    )
    created_at: datetime = Field(
        default_factory=lambda: datetime.now(timezone.utc),
        description="Creation timestamp (UTC).",
    )
    updated_at: datetime = Field(
        default_factory=lambda: datetime.now(timezone.utc),
        description="Last modification timestamp (UTC).",
    )

    embedding: List[float] | None = Field(
        default=None,
        exclude=True,
        description="Optional cached embedding for retrieval tasks.",
    )

    model_config = ConfigDict(use_enum_values=True)


class ReflectionReport(BaseModel):
    """Structure generated by the reflector to describe issues and new insights."""

    question: str
    answer: str
    evidence_refs: List[str] = Field(default_factory=list)
    issues: List[ReflectionReport.Issue] = Field(
        default_factory=list,
        description="Identified issues such as omissions or contradictions.",
    )
    proposed_insights: List[ReflectionReport.ProposedInsight] = Field(
        default_factory=list,
        description="Candidate deltas proposed by the reflector.",
    )

    class Issue(BaseModel):
        """Information about a detected issue during reflection."""

        type: str
        explanation: str
        severity: int = Field(ge=1, le=5)

        model_config = ConfigDict(extra="forbid")

    class ProposedInsight(BaseModel):
        """Structure describing a potential new delta."""

        topic: str
        guideline: str
        conditions: List[str] = Field(default_factory=list)
        evidence: List[str] = Field(default_factory=list)
        tags: List[str] = Field(default_factory=list)
        confidence: float = Field(ge=0.0, le=1.0)

        model_config = ConfigDict(extra="forbid")

    model_config = ConfigDict(extra="allow")


class ACETConfig(BaseModel):
    """Runtime configuration for the ACET engine."""

    token_budget: int = Field(
        default=800,
        description="Maximum number of tokens allocated for context injection.",
    )
    batch_size: int = Field(
        default=1,
        ge=1,
        description="Number of deltas processed in a single batch.",
    )
    max_epochs: int = Field(
        default=5,
        ge=1,
        description="Maximum offline training epochs.",
    )
    max_reflector_iterations: int = Field(
        default=5,
        ge=1,
        description="Maximum iterations when refining reflection reports.",
    )

    similarity_weight: float = Field(
        default=0.45,
        ge=0.0,
        le=1.0,
        description="Weight applied to semantic similarity during ranking.",
    )
    recency_weight: float = Field(
        default=0.25,
        ge=0.0,
        le=1.0,
        description="Weight applied to recency during ranking.",
    )
    usage_weight: float = Field(
        default=0.20,
        ge=0.0,
        le=1.0,
        description="Weight applied to usage statistics during ranking.",
    )
    risk_penalty_weight: float = Field(
        default=0.10,
        ge=0.0,
        le=1.0,
        description="Penalty weight applied to risky deltas during ranking.",
    )

    min_confidence: float = Field(
        default=0.55,
        ge=0.0,
        le=1.0,
        description="Minimum confidence threshold for accepting new deltas.",
    )
    dedup_threshold: float = Field(
        default=0.90,
        ge=0.0,
        le=1.0,
        description="Similarity threshold used for deduplication.",
    )

    reflection_sample_rate: float = Field(
        default=0.3,
        ge=0.0,
        le=1.0,
        description="Probability of running reflection on a generation.",
    )

