"""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

__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 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)
    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 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 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.payment = PaymentContract(self._state, cluster, starting_credit_lamports)
        self.settlement = SettlementContract(
            self._state,
            cluster,
            cost_per_layer_token_lamport,
        )


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