"""Solana contract boundary for the Distributed Inference Network.

The prototype uses deterministic local state with Solana-shaped wrapper classes.
ADR-0007 keeps the public Python boundary stable so real Solana programs can
replace this adapter later without changing gateway or tracker call sites.
"""

from dataclasses import dataclass, field
import json
import urllib.request

__version__ = "0.1.0"


@dataclass(frozen=True)
class RegistryWallet:
    """Stake, strike, and ban state for a node operator wallet."""

    stake_balance: int = 0
    strike_count: int = 0
    banned: bool = False


@dataclass(frozen=True)
class ApiKeyBalance:
    """Client API key payment account balance."""

    lamports: int = 0
    usdc_micro: int = 0


@dataclass(frozen=True)
class ComputeAttribution:
    """On-chain work attribution recorded by the gateway after inference."""

    session_id: str
    api_key: str
    node_wallet: str
    layer_start: int
    layer_end: int
    tokens: int
    speed_score: float = 1.0
    settled_epoch: int | None = None

    @property
    def layer_count(self) -> int:
        return self.layer_end - self.layer_start + 1

    @property
    def work_units(self) -> int:
        return self.layer_count * self.tokens


@dataclass(frozen=True)
class ValidationEvent:
    """Completed inference data consumed by fraud validators."""

    index: int
    session_id: str
    model: str
    messages: list[dict]
    observed_output: str
    route_nodes: list[dict]


@dataclass(frozen=True)
class SlashReceipt:
    """Local slash transaction receipt."""

    signature: str
    cluster: str
    wallet_address: str
    slash_amount: int
    stake_before: int
    stake_after: int
    strike_count: int
    banned: bool
    reason: str


@dataclass(frozen=True)
class SettlementResult:
    """Summary returned by an epoch settlement transaction."""

    epoch: int
    client_debits: dict[str, int]
    node_rewards: dict[str, int]
    validator_rewards: dict[str, int]


@dataclass
class _LocalContractState:
    registry: dict[str, RegistryWallet] = field(default_factory=dict)
    api_keys: dict[str, ApiKeyBalance] = field(default_factory=dict)
    attributions: list[ComputeAttribution] = field(default_factory=list)
    validation_events: list[ValidationEvent] = field(default_factory=list)
    slash_receipts: list[SlashReceipt] = field(default_factory=list)
    token_balances: dict[str, int] = field(default_factory=dict)


class RegistryContract:
    """Registry wrapper for node stake, strikes, and bans."""

    def __init__(self, state: _LocalContractState, cluster: str) -> None:
        self._state = state
        self._cluster = cluster

    def submit_stake(self, wallet_address: str, amount: int) -> dict:
        if amount <= 0:
            raise ValueError("stake amount must be positive")
        current = self.get_wallet(wallet_address)
        self._state.registry[wallet_address] = RegistryWallet(
            stake_balance=current.stake_balance + amount,
            strike_count=current.strike_count,
            banned=current.banned,
        )
        return {"signature": f"local-stake-{wallet_address}", "cluster": self._cluster}

    def get_wallet(self, wallet_address: str) -> RegistryWallet:
        return self._state.registry.get(wallet_address, RegistryWallet())

    def record_strike(self, wallet_address: str, ban_threshold: int = 3) -> dict:
        current = self.get_wallet(wallet_address)
        strike_count = current.strike_count + 1
        self._state.registry[wallet_address] = RegistryWallet(
            stake_balance=current.stake_balance,
            strike_count=strike_count,
            banned=current.banned or strike_count >= ban_threshold,
        )
        return {"signature": f"local-strike-{wallet_address}", "cluster": self._cluster}

    def submit_slash_proof(
        self,
        *,
        wallet_address: str,
        slash_amount: int,
        strike_threshold: int = 3,
        reason: str = "fraudulent inference output",
        webhook_url: str | None = None,
    ) -> SlashReceipt:
        if not wallet_address:
            raise ValueError("wallet_address is required")
        if slash_amount <= 0:
            raise ValueError("slash_amount must be positive")
        if strike_threshold <= 0:
            raise ValueError("strike_threshold must be positive")

        current = self.get_wallet(wallet_address)
        stake_after = max(0, current.stake_balance - slash_amount)
        strike_count = current.strike_count + 1
        banned = current.banned or strike_count >= strike_threshold
        self._state.registry[wallet_address] = RegistryWallet(
            stake_balance=stake_after,
            strike_count=strike_count,
            banned=banned,
        )
        receipt = SlashReceipt(
            signature=f"local-slash-{wallet_address}-{len(self._state.slash_receipts) + 1}",
            cluster=self._cluster,
            wallet_address=wallet_address,
            slash_amount=slash_amount,
            stake_before=current.stake_balance,
            stake_after=stake_after,
            strike_count=strike_count,
            banned=banned,
            reason=reason,
        )
        self._state.slash_receipts.append(receipt)
        _notify_slash(receipt, webhook_url)
        return receipt

    def ban_wallet(self, wallet_address: str) -> dict:
        current = self.get_wallet(wallet_address)
        self._state.registry[wallet_address] = RegistryWallet(
            stake_balance=current.stake_balance,
            strike_count=current.strike_count,
            banned=True,
        )
        return {"signature": f"local-ban-{wallet_address}", "cluster": self._cluster}

    def has_minimum_stake(self, wallet_address: str | None, minimum_stake: int) -> bool:
        if not wallet_address:
            return False
        wallet = self.get_wallet(wallet_address)
        return wallet.stake_balance >= minimum_stake


class ValidationContract:
    """Validation event log consumed by the optimistic fraud detector."""

    def __init__(self, state: _LocalContractState) -> None:
        self._state = state

    def record_completed_inference(
        self,
        *,
        session_id: str,
        model: str,
        messages: list[dict],
        observed_output: str,
        route_nodes: list[dict],
    ) -> ValidationEvent:
        if not session_id:
            raise ValueError("session_id is required")
        if not model:
            raise ValueError("model is required")
        if not isinstance(messages, list):
            raise ValueError("messages must be a list")
        if not isinstance(observed_output, str):
            raise ValueError("observed_output must be a string")
        if not isinstance(route_nodes, list):
            raise ValueError("route_nodes must be a list")
        event = ValidationEvent(
            index=len(self._state.validation_events),
            session_id=session_id,
            model=model,
            messages=[dict(message) for message in messages],
            observed_output=observed_output,
            route_nodes=[dict(node) for node in route_nodes],
        )
        self._state.validation_events.append(event)
        return event

    def list_completed_inferences(self, *, after_index: int = -1) -> list[ValidationEvent]:
        return [
            event for event in self._state.validation_events
            if event.index > after_index
        ]


class PaymentContract:
    """Payment wrapper for funded API keys and compute attribution."""

    def __init__(self, state: _LocalContractState, cluster: str, starting_credit_lamports: int) -> None:
        self._state = state
        self._cluster = cluster
        self._starting_credit_lamports = starting_credit_lamports

    def fund_api_key(
        self,
        api_key: str,
        *,
        lamports: int = 0,
        usdc_micro: int = 0,
    ) -> dict:
        if lamports < 0 or usdc_micro < 0:
            raise ValueError("funding amounts must be non-negative")
        if lamports == 0 and usdc_micro == 0:
            raise ValueError("funding amount must be positive")
        current = self.get_balance(api_key)
        self._state.api_keys[api_key] = ApiKeyBalance(
            lamports=current.lamports + lamports,
            usdc_micro=current.usdc_micro + usdc_micro,
        )
        return {"signature": f"local-fund-{api_key}", "cluster": self._cluster}

    def get_balance(self, api_key: str) -> ApiKeyBalance:
        return self._state.api_keys.get(
            api_key,
            ApiKeyBalance(lamports=self._starting_credit_lamports),
        )

    def record_attribution(
        self,
        *,
        session_id: str,
        api_key: str,
        node_wallet: str,
        layer_start: int,
        layer_end: int,
        tokens: int,
        speed_score: float = 1.0,
    ) -> dict:
        if not api_key:
            raise ValueError("api_key is required")
        if not node_wallet:
            raise ValueError("node_wallet is required")
        if layer_start < 0 or layer_end < layer_start:
            raise ValueError("layer range must be non-negative and ordered")
        if tokens <= 0:
            raise ValueError("tokens must be positive")
        if speed_score <= 0:
            raise ValueError("speed_score must be positive")
        self._state.attributions.append(ComputeAttribution(
            session_id=session_id,
            api_key=api_key,
            node_wallet=node_wallet,
            layer_start=layer_start,
            layer_end=layer_end,
            tokens=tokens,
            speed_score=speed_score,
        ))
        return {"signature": f"local-attribution-{session_id}", "cluster": self._cluster}

    def list_attributions(self, *, include_settled: bool = True) -> list[ComputeAttribution]:
        if include_settled:
            return list(self._state.attributions)
        return [a for a in self._state.attributions if a.settled_epoch is None]


class SettlementContract:
    """Settlement wrapper that debits clients and credits token rewards."""

    def __init__(
        self,
        state: _LocalContractState,
        cluster: str,
        cost_per_layer_token_lamport: int,
    ) -> None:
        self._state = state
        self._cluster = cluster
        self._cost_per_layer_token_lamport = cost_per_layer_token_lamport

    def settle_epoch(
        self,
        *,
        epoch: int,
        validator_wallet: str | None = None,
        validator_reward_share_bps: int = 0,
    ) -> SettlementResult:
        if epoch < 0:
            raise ValueError("epoch must be non-negative")
        if not 0 <= validator_reward_share_bps <= 10_000:
            raise ValueError("validator_reward_share_bps must be between 0 and 10000")
        unsettled = [a for a in self._state.attributions if a.settled_epoch is None]

        client_debits: dict[str, int] = {}
        work_by_node: dict[str, int] = {}
        for attribution in unsettled:
            debit = attribution.work_units * self._cost_per_layer_token_lamport
            weighted_work = int(attribution.work_units * attribution.speed_score)
            client_debits[attribution.api_key] = client_debits.get(attribution.api_key, 0) + debit
            work_by_node[attribution.node_wallet] = (
                work_by_node.get(attribution.node_wallet, 0) + weighted_work
            )

        for api_key, debit in client_debits.items():
            current = self._state.api_keys.get(api_key, ApiKeyBalance())
            if current.lamports < debit:
                raise ValueError(f"insufficient SOL balance for API key: {api_key}")
            self._state.api_keys[api_key] = ApiKeyBalance(
                lamports=current.lamports - debit,
                usdc_micro=current.usdc_micro,
            )

        total_weighted_work = sum(work_by_node.values())
        raw_total_work = sum(a.work_units for a in unsettled)
        validator_rewards: dict[str, int] = {}
        validator_reward_total = (raw_total_work * validator_reward_share_bps) // 10_000
        if validator_wallet and validator_reward_total:
            validator_rewards[validator_wallet] = validator_reward_total

        node_reward_pool = raw_total_work - validator_reward_total
        node_rewards: dict[str, int] = {}
        for node_wallet, weighted_work_units in work_by_node.items():
            reward = 0 if total_weighted_work == 0 else (node_reward_pool * weighted_work_units) // total_weighted_work
            node_rewards[node_wallet] = reward
            self._state.token_balances[node_wallet] = (
                self._state.token_balances.get(node_wallet, 0) + reward
            )

        for wallet, reward in validator_rewards.items():
            self._state.token_balances[wallet] = self._state.token_balances.get(wallet, 0) + reward

        self._state.attributions = [
            ComputeAttribution(
                session_id=a.session_id,
                api_key=a.api_key,
                node_wallet=a.node_wallet,
                layer_start=a.layer_start,
                layer_end=a.layer_end,
                tokens=a.tokens,
                speed_score=a.speed_score,
                settled_epoch=epoch,
            )
            if a.settled_epoch is None else a
            for a in self._state.attributions
        ]
        return SettlementResult(
            epoch=epoch,
            client_debits=client_debits,
            node_rewards=node_rewards,
            validator_rewards=validator_rewards,
        )

    def get_token_balance(self, wallet_address: str) -> int:
        return self._state.token_balances.get(wallet_address, 0)

    def deployment_plan(self) -> dict:
        """Return the configured manual testnet deployment targets."""
        return {
            "cluster": self._cluster,
            "programs": ["registry", "payment", "settlement"],
        }


class LocalSolanaContracts:
    """Facade that exposes all three contract wrappers over local validator state."""

    def __init__(
        self,
        *,
        cluster: str = "local-test-validator",
        cost_per_layer_token_lamport: int = 1,
        starting_credit_lamports: int = 1_000,
    ) -> None:
        if cost_per_layer_token_lamport <= 0:
            raise ValueError("cost_per_layer_token_lamport must be positive")
        if starting_credit_lamports < 0:
            raise ValueError("starting_credit_lamports must be non-negative")
        self.cluster = cluster
        self._state = _LocalContractState()
        self.registry = RegistryContract(self._state, cluster)
        self.validation = ValidationContract(self._state)
        self.payment = PaymentContract(self._state, cluster, starting_credit_lamports)
        self.settlement = SettlementContract(
            self._state,
            cluster,
            cost_per_layer_token_lamport,
        )


def _notify_slash(receipt: SlashReceipt, webhook_url: str | None) -> None:
    message = (
        f"WARNING: node {receipt.wallet_address} slashed "
        f"by {receipt.slash_amount} lamports; strikes={receipt.strike_count}; "
        f"banned={receipt.banned}; reason={receipt.reason}"
    )
    print(message, flush=True)
    if not webhook_url:
        return
    payload = json.dumps({
        "event": "node_slashed",
        "wallet_address": receipt.wallet_address,
        "slash_amount": receipt.slash_amount,
        "strike_count": receipt.strike_count,
        "banned": receipt.banned,
        "reason": receipt.reason,
    }).encode()
    req = urllib.request.Request(
        webhook_url,
        data=payload,
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    try:
        urllib.request.urlopen(req, timeout=2).read()
    except OSError:
        pass


__all__ = [
    "ApiKeyBalance",
    "ComputeAttribution",
    "LocalSolanaContracts",
    "PaymentContract",
    "RegistryContract",
    "RegistryWallet",
    "SettlementContract",
    "SettlementResult",
    "SlashReceipt",
    "ValidationContract",
    "ValidationEvent",
]