from __future__ import annotations

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


def _stable_json(payload: Dict[str, Any]) -> str:
    def normalize(value: Any) -> Any:
        if isinstance(value, dict):
            return {str(key): normalize(val) for key, val in sorted(value.items(), key=lambda item: _stable_key(item[0]))}
        if isinstance(value, (list, tuple)):
            return [normalize(item) for item in value]
        return value

    return json.dumps(normalize(payload), sort_keys=True, separators=(",", ":"), default=str)


def _stable_key(item: Any) -> str:
    return repr(item)


@dataclass
class LocalArbProblem:
    id: str
    venue: str
    assets: List[str]
    target_misprice: float
    max_exposure: float
    latency_budget: float

    def to_dict(self) -> Dict[str, Any]:
        return {
            "id": self.id,
            "venue": self.venue,
            "assets": self.assets,
            "target_misprice": self.target_misprice,
            "max_exposure": self.max_exposure,
            "latency_budget": self.latency_budget,
        }

    def fingerprint(self) -> str:
        return hashlib.sha256(_stable_json(self.to_dict()).encode("utf-8")).hexdigest()


@dataclass
class SharedSignals:
    version: int
    price_delta_by_asset: Dict[str, float] = field(default_factory=dict)
    cross_corr: Dict[Any, float] = field(default_factory=dict)
    liquidity_estimates: Dict[str, float] = field(default_factory=dict)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "version": self.version,
            "price_delta_by_asset": dict(self.price_delta_by_asset),
            "cross_corr": dict(self.cross_corr),
            "liquidity_estimates": dict(self.liquidity_estimates),
        }

    def fingerprint(self) -> str:
        return hashlib.sha256(_stable_json(self.to_dict()).encode("utf-8")).hexdigest()

    def merge(self, other: "SharedSignals") -> "SharedSignals":
        merged = SharedSignals(version=max(self.version, other.version))
        keys = sorted(set(self.price_delta_by_asset) | set(other.price_delta_by_asset), key=_stable_key)
        for k in keys:
            a = self.price_delta_by_asset.get(k, 0.0)
            b = other.price_delta_by_asset.get(k, 0.0)
            merged.price_delta_by_asset[k] = a if abs(a) >= abs(b) else b
        for k in sorted(set(self.cross_corr) | set(other.cross_corr), key=_stable_key):
            merged.cross_corr[k] = (self.cross_corr.get(k, 0.0) + other.cross_corr.get(k, 0.0)) / 2.0
        for k in sorted(set(self.liquidity_estimates) | set(other.liquidity_estimates), key=_stable_key):
            merged.liquidity_estimates[k] = (
                self.liquidity_estimates.get(k, 0.0) + other.liquidity_estimates.get(k, 0.0)
            ) / 2.0
        return merged


@dataclass
class PlanDelta:
    delta_actions: List[Dict[str, Any]]
    timestamp: float
    contract_id: str
    signature: str = ""

    def __post_init__(self) -> None:
        if not self.signature:
            self.signature = self.fingerprint()

    def to_dict(self) -> Dict[str, Any]:
        return {
            "delta_actions": self.delta_actions,
            "timestamp": self.timestamp,
            "contract_id": self.contract_id,
            "signature": self.signature,
        }

    def fingerprint(self) -> str:
        payload = {
            "contract_id": self.contract_id,
            "delta_actions": self.delta_actions,
            "timestamp": self.timestamp,
        }
        return hashlib.sha256(_stable_json(payload).encode("utf-8")).hexdigest()


@dataclass
class DualVariables:
    shadow_price: float = 0.0

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


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

    def log(self, msg: str) -> None:
        prev = self.hashes[-1] if self.hashes else ""
        entry_hash = hashlib.sha256(f"{prev}|{msg}".encode("utf-8")).hexdigest()
        self.entries.append(msg)
        self.hashes.append(entry_hash)

    def root_hash(self) -> str:
        return self.hashes[-1] if self.hashes else ""


@dataclass
class PrivacyBudget:
    leakage_budget: float

    def remaining(self) -> float:
        return self.leakage_budget

    def can_consume(self, amount: float) -> bool:
        return amount <= self.leakage_budget

    def consume(self, amount: float) -> None:
        self.leakage_budget = max(0.0, self.leakage_budget - amount)