from __future__ import annotations

from datetime import datetime, timezone
from typing import Any

from pydantic import BaseModel, Field, field_validator, model_validator


class LocalProblem(BaseModel):
    robot_id: str
    target_pos: float
    initial_pos: float = 0.0
    current_pos: float | None = None
    dual_var: float = 0.0
    max_step: float = 15.0
    energy_budget: float = 100.0
    weight: float = 1.0
    version: int = 0
    metadata: dict[str, Any] = Field(default_factory=dict)

    @field_validator("max_step", "energy_budget", "weight")
    @classmethod
    def _positive(cls, value: float) -> float:
        if value <= 0:
            raise ValueError("must be positive")
        return value

    @model_validator(mode="after")
    def _default_current_pos(self) -> "LocalProblem":
        if self.current_pos is None:
            self.current_pos = self.initial_pos
        return self

    @property
    def travel_distance(self) -> float:
        current_pos = self.current_pos if self.current_pos is not None else self.initial_pos
        return abs(current_pos - self.initial_pos)

    def apply_updates(self, updates: dict[str, Any]) -> None:
        for key, value in updates.items():
            if hasattr(self, key):
                setattr(self, key, value)


class SharedVariables(BaseModel):
    version: int = 0
    global_consensus: float = 0.0
    delta_clock: int = 0


class DualVariables(BaseModel):
    values: dict[str, float] = Field(default_factory=dict)


class PlanDelta(BaseModel):
    author: str
    contract_id: str
    sequence: int
    timestamp: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
    target_robot_id: str | None = None
    updates: dict[str, Any] = Field(default_factory=dict)

    @field_validator("sequence")
    @classmethod
    def _non_negative(cls, value: int) -> int:
        if value < 0:
            raise ValueError("sequence must be non-negative")
        return value


class SafetyPolicy(BaseModel):
    max_travel_distance: float = 15.0
    min_separation: float = 1.0
    energy_budget: float = 100.0
    bounded_staleness: int = 2

    @field_validator("max_travel_distance", "min_separation", "energy_budget")
    @classmethod
    def _non_negative(cls, value: float) -> float:
        if value < 0:
            raise ValueError("must be non-negative")
        return value

    @field_validator("bounded_staleness")
    @classmethod
    def _bounded_staleness(cls, value: int) -> int:
        if value < 0:
            raise ValueError("must be non-negative")
        return value

    def verify(self, start_pos: float, new_pos: float) -> bool:
        return abs(new_pos - start_pos) <= self.max_travel_distance


class AuditLogEntry(BaseModel):
    iteration: int
    shared_version: int
    consensus: float
    primal_residual: float
    dual_residual: float
    applied_deltas: list[str] = Field(default_factory=list)
    notes: list[str] = Field(default_factory=list)
    created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))


class ConvergenceCertificate(BaseModel):
    converged: bool
    iterations: int
    epsilon: float
    final_primal_residual: float
    final_dual_residual: float


class SwarmSolution(BaseModel):
    consensus: float
    robots: list[LocalProblem]
    shared: SharedVariables
    audit_log: list[AuditLogEntry]
    certificate: ConvergenceCertificate