from __future__ import annotations

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


@dataclass
class LocalProblem:
    id: str
    domain: str
    assets: List[str]
    objective: str
    constraints: Dict[str, Any] = field(default_factory=dict)


@dataclass
class SharedVariables:
    forecasts: Dict[str, float] = field(default_factory=dict)
    priors: Dict[str, float] = field(default_factory=dict)
    version: int = 0


@dataclass
class PlanDelta:
    delta: Dict[str, Any]
    timestamp: datetime
    author: str
    contract_id: str
    signature: Optional[str] = None
    # Optional time-round metadata to support deterministic offline replay
    time_rounds: Optional["TimeRounds"] = None


@dataclass
class DualVariables:
    lambda_latency: float = 0.0
    lambda_fees: float = 0.0


@dataclass
class PrivacyBudget:
    budget: float
    expiry: Optional[datetime] = None


@dataclass
class AuditLog:
    entry: str
    signer: str
    timestamp: datetime
    contract_id: str
    version: int


@dataclass
class GraphOfContractsRegistry:
    registry: Dict[str, str] = field(default_factory=dict)

    def register(self, adapter_id: str, contract_version: str) -> None:
        self.registry[adapter_id] = contract_version

    def get_version(self, adapter_id: str) -> Optional[str]:
        return self.registry.get(adapter_id)


@dataclass
class TimeRounds:
    """Monoid-like container describing discrete offline replay rounds."""
    round_id: str
    start_ts: float
    end_ts: float


@dataclass
class GoCEntry:
    adapter_id: str
    supported_domains: List[str]
    contract_version: str


@dataclass
class GoCRegistry:
    entries: List[GoCEntry] = field(default_factory=list)

    def register(self, entry: GoCEntry) -> None:
        self.entries.append(entry)

    def get_contract_version(self, adapter_id: str) -> Optional[str]:
        for e in self.entries:
            if e.adapter_id == adapter_id:
                return e.contract_version
        return None


__all__ = [
    "LocalProblem",
    "SharedVariables",
    "PlanDelta",
    "DualVariables",
    "PrivacyBudget",
    "AuditLog",
    "GraphOfContractsRegistry",
    "TimeRounds",
    "GoCEntry",
    "GoCRegistry",
]