from __future__ import annotations

import json
from dataclasses import dataclass, field
import hashlib
from typing import List, Dict, Any, Optional


def _to_json(obj: Any) -> Any:
    if hasattr(obj, "to_json"):
        return obj.to_json()
    if isinstance(obj, list):
        return [_to_json(v) for v in obj]
    if isinstance(obj, dict):
        return {k: _to_json(v) for k, v in obj.items()}
    return obj


def _from_json(cls, data: Any):
    if hasattr(cls, "from_json"):
        return cls.from_json(data)
    return data


class LocalArbProblem:
    def __init__(
        self,
        venue_id: str = "DEFAULT_VENUE",
        asset_pair=None,
        target_mispricing: float = 0.0,
        liquidity_budget: float = 0.0,
        latency_budget_ms: int | None = None,
        latency_budget: float | None = None,
    ):
        self.venue_id = venue_id
        self.asset_pair = list(asset_pair) if asset_pair is not None else []
        self.target_mispricing = float(target_mispricing)
        self.liquidity_budget = float(liquidity_budget)
        self.latency_budget_ms = latency_budget_ms
        self.latency_budget = latency_budget
        if self.latency_budget_ms is None and self.latency_budget is not None:
            self.latency_budget_ms = int(self.latency_budget * 1000)

    def to_json(self) -> Dict[str, Any]:
        return {
            "venue_id": self.venue_id,
            "asset_pair": self.asset_pair,
            "target_mispricing": self.target_mispricing,
            "liquidity_budget": self.liquidity_budget,
            "latency_budget_ms": self.latency_budget_ms,
        }

    @classmethod
    def from_json(cls, data: Dict[str, Any]) -> "LocalArbProblem":
        return cls(
            venue_id=data.get("venue_id", "DEFAULT_VENUE"),
            asset_pair=data.get("asset_pair", []),
            target_mispricing=data.get("target_mispricing", 0.0),
            liquidity_budget=data.get("liquidity_budget", 0.0),
            latency_budget_ms=data.get("latency_budget_ms"),
        )

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, LocalArbProblem):
            return False
        return (
            self.venue_id == other.venue_id
            and self.asset_pair == other.asset_pair
            and self.target_mispricing == other.target_mispricing
            and self.liquidity_budget == other.liquidity_budget
            and self.latency_budget_ms == other.latency_budget_ms
        )


class SharedSignals:
    def __init__(self, deltas: list[float], cross_venue_corr: float, liquidity_availability: float, latency_proxy: float):
        self.deltas = deltas
        self.cross_venue_corr = cross_venue_corr
        self.liquidity_availability = liquidity_availability
        self.latency_proxy = latency_proxy

    def to_json(self) -> Dict[str, Any]:
        return {
            "deltas": self.deltas,
            "cross_venue_corr": self.cross_venue_corr,
            "liquidity_availability": self.liquidity_availability,
            "latency_proxy": self.latency_proxy,
        }

    @classmethod
    def from_json(cls, data: Dict[str, Any]) -> "SharedSignals":
        return cls(
            deltas=data.get("deltas", []),
            cross_venue_corr=data.get("cross_venue_corr", 0.0),
            liquidity_availability=data.get("liquidity_availability", 0.0),
            latency_proxy=data.get("latency_proxy", 0.0),
        )


@dataclass
class PlanDelta:
    delta: List[Dict[str, Any]]
    timestamp: float
    author: str
    signature: Optional[str] = None

    def to_json(self) -> Dict[str, Any]:
        return {
            "delta": self.delta,
            "timestamp": self.timestamp,
            "author": self.author,
            "signature": self.signature,
        }

    @classmethod
    def from_json(cls, data: Dict[str, Any]) -> "PlanDelta":
        return cls(
            delta=data["delta"],
            timestamp=data["timestamp"],
            author=data["author"],
            signature=data.get("signature"),
        )


class DualVariables:
    def __init__(self, shadow_prices: Dict[str, float] | None = None, venue_shadows: Dict[str, float] | None = None):
        if shadow_prices is not None:
            self.venue_shadows = dict(shadow_prices)
        elif venue_shadows is not None:
            self.venue_shadows = dict(venue_shadows)
        else:
            self.venue_shadows = {}

    def to_json(self) -> Dict[str, Any]:
        return {"venue_shadows": self.venue_shadows}

    @classmethod
    def from_json(cls, data: Dict[str, Any]) -> "DualVariables":
        return cls(venue_shadows=data.get("venue_shadows", {}))

    # Compatibility: provide shadow_prices alias used by existing IR mapping
    @property
    def shadow_prices(self) -> Dict[str, float]:
        return self.venue_shadows


@dataclass
class PrivacyBudget:
    def __init__(self, total_budget: float | None = None, spent: float = 0.0, budgets: Dict[str, float] | None = None):
        if budgets is not None:
            self.total_budget = budgets.get("total", budgets.get("signals", 0.0))
            self.spent = budgets.get("spent", 0.0)
        else:
            self.total_budget = total_budget if total_budget is not None else 0.0
            self.spent = spent

    def to_json(self) -> Dict[str, Any]:
        return {"total_budget": self.total_budget, "spent": self.spent}

    @classmethod
    def from_json(cls, data: Dict[str, Any]) -> "PrivacyBudget":
        return cls(total_budget=data.get("total_budget", 0.0), spent=data.get("spent", 0.0))

    @property
    def budgets(self) -> Dict[str, float]:
        # Compatibility: expose a single "signals" budget for the IR consumer
        return {"signals": self.total_budget}


@dataclass
class AuditLogEntry:
    def __init__(self, ts: float, event: str, signer: Optional[str] = None, details: Dict[str, Any] = None, signature: Optional[str] = None):
        self.ts = ts
        self.event = event
        # Support both 'signer' and legacy 'signature' naming
        self.signer = signer if signer is not None else signature
        self.details = details if details is not None else {}

    def to_json(self) -> Dict[str, Any]:
        return {"event": self.event, "ts": self.ts, "signer": self.signer, "details": self.details}

    @classmethod
    def from_json(cls, data: Dict[str, Any]) -> "AuditLogEntry":
        return cls(ts=data["ts"], event=data["event"], signer=data.get("signer"), details=data.get("details", {}))

    @property
    def signature(self) -> str | None:
        return self.signer


@dataclass
class AuditLog:
    entries: List[AuditLogEntry] = field(default_factory=list)

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

    def to_json(self) -> Dict[str, Any]:
        serialized = []
        for e in self.entries:
            if hasattr(e, "to_json"):
                serialized.append(e.to_json())
            else:
                serialized.append(e)
        return {"entries": serialized}

    @classmethod
    def from_json(cls, data: Dict[str, Any]) -> "AuditLog":
        entries = [AuditLogEntry.from_json(e) for e in data.get("entries", [])]
        return cls(entries=entries)


@dataclass
class TimeRounds:
    rounds: List[int] = field(default_factory=list)

    def to_json(self) -> Dict[str, Any]:
        return {"rounds": self.rounds}

    @classmethod
    def from_json(cls, data: Dict[str, Any]) -> "TimeRounds":
        return cls(rounds=data.get("rounds", []))


@dataclass
class GoCRegistry:
    adapters: Dict[str, str] = field(default_factory=dict)  # adapter_name -> contract_version

    def to_json(self) -> Dict[str, Any]:
        return {"adapters": self.adapters}

    @classmethod
    def from_json(cls, data: Dict[str, Any]) -> "GoCRegistry":
        return cls(adapters=data.get("adapters", {}))


@dataclass
class DeltaTrace:
    """Deterministic replay trace for a PlanDelta per venue.

    This is a lightweight hash- chained log entry capturing the delta payload
    and its provenance. It enables auditable backtesting and partition reconciliation
    without exposing raw delta details beyond the payload envelope.
    """
    delta_hash: str
    parent_hash: str | None
    timestamp: float
    signer: str
    payload: Dict[str, Any]

    def to_json(self) -> Dict[str, Any]:
        return {
            "delta_hash": self.delta_hash,
            "parent_hash": self.parent_hash,
            "timestamp": self.timestamp,
            "signer": self.signer,
            "payload": self.payload,
        }

    @classmethod
    def from_json(cls, data: Dict[str, Any]) -> "DeltaTrace":
        return cls(
            delta_hash=data["delta_hash"],
            parent_hash=data.get("parent_hash"),
            timestamp=data["timestamp"],
            signer=data["signer"],
            payload=data["payload"],
        )

    @staticmethod
    def compute_hash(payload: Dict[str, Any]) -> str:
        # Stable hash of the JSON payload
        payload_bytes = json.dumps(payload, sort_keys=True, separators=(",", ":")).encode("utf-8")
        return hashlib.sha256(payload_bytes).hexdigest()