"""Off-chain USDT billing ledger (ADR-0015, US-031).

Tracks client API-key balances, node pending balances, and the protocol cut.
All mutations are expressed as append-only billing events with unique ids so
the ledger converges when events are gossiped across the tracker hive: every
event is a commutative balance delta, and each tracker applies an event at
most once (dedupe by event id).

No Solana calls live here — on-chain deposit/settlement adapters (US-032/033)
consume this ledger through `snapshot()` and `settle_node_payout()`.
"""

from __future__ import annotations

import json
import sqlite3
import threading
import time
import uuid

DEFAULT_PRICE_PER_1K_TOKENS = 0.02  # USDT
DEFAULT_STARTING_CREDIT = 0.0  # USDT; accounts must be funded before inference
DEFAULT_BILLING_DB_PATH = "billing.sqlite"
NODE_REVENUE_SHARE = 0.90  # nodes 90% / protocol cut 10% (ADR-0015)


def _normalize_rates(value: "float | tuple[float, float] | dict") -> tuple[float, float]:
    """Coerce a price spec into an (input_per_1k, output_per_1k) pair."""
    if isinstance(value, dict):
        base = value.get("price")
        inp = value.get("input", base)
        out = value.get("output", base)
        return (float(inp), float(out))
    if isinstance(value, (tuple, list)):
        return (float(value[0]), float(value[1]))
    return (float(value), float(value))


class BillingLedger:
    """Thread-safe USDT ledger with SQLite persistence and event replication."""

    SAVE_INTERVAL = 10.0

    def __init__(
        self,
        db_path: str | None = None,
        *,
        prices: dict[str, float | tuple[float, float]] | None = None,
        default_price_per_1k: float = DEFAULT_PRICE_PER_1K_TOKENS,
        starting_credit: float = DEFAULT_STARTING_CREDIT,
        node_share: float = NODE_REVENUE_SHARE,
    ) -> None:
        self._db_path = db_path
        # US-045: per-model (input_per_1k, output_per_1k). A bare float in
        # ``prices`` sets both rates (legacy single-rate models).
        self._prices: dict[str, tuple[float, float]] = {
            model: _normalize_rates(value) for model, value in (prices or {}).items()
        }
        self._default_price_per_1k = default_price_per_1k
        self._starting_credit = starting_credit
        self._node_share = node_share
        self._lock = threading.Lock()
        self._client_balances: dict[str, float] = {}
        self._node_pending: dict[str, float] = {}
        self._wallet_bindings: dict[str, str] = {}  # client wallet pubkey -> api_key
        self._protocol_cut: float = 0.0
        self._pending_since: dict[str, float] = {}  # wallet -> ts pending became > 0
        self._last_payout_ts: dict[str, float] = {}
        # settlement id -> {"payouts": [(wallet, amount)], "signature": str|None, "ts": float}
        self._settlements: dict[str, dict] = {}
        self._seen_event_ids: set[str] = set()
        self._event_log: list[dict] = []  # full log, order of local application
        self._dirty = False
        if self._db_path:
            self._init_db()
            self._load_from_db()

    # ---- pricing ----

    def price_for(self, model: str) -> float:
        """Blended (average) per-1k rate — kept for estimators and history logs."""
        rates = self._prices.get(model)
        if rates is None:
            return self._default_price_per_1k
        return (rates[0] + rates[1]) / 2.0

    def prices_for(self, model: str) -> tuple[float, float]:
        """(input_per_1k, output_per_1k) for a model (US-045)."""
        return self._prices.get(model, (self._default_price_per_1k, self._default_price_per_1k))

    def set_price(self, model: str, price_per_1k: float) -> None:
        """Legacy single-rate setter — applies the same rate to input and output."""
        self.set_prices(model, price_per_1k, price_per_1k)

    def set_prices(self, model: str, input_per_1k: float, output_per_1k: float) -> None:
        with self._lock:
            self._prices[model] = (float(input_per_1k), float(output_per_1k))

    # ---- local operations (create + apply + log an event) ----

    def ensure_client(self, api_key: str) -> float:
        """Return the client's balance without granting implicit credit."""
        with self._lock:
            if api_key not in self._client_balances and self._starting_credit > 0.0:
                self._apply_locked({
                    "id": f"credit-{uuid.uuid4().hex}",
                    "type": "credit",
                    "api_key": api_key,
                    "amount": self._starting_credit,
                    "ts": time.time(),
                    "note": "caller-credit",
                })
            return self._client_balances.get(api_key, 0.0)

    def has_funds(self, api_key: str) -> bool:
        return self.ensure_client(api_key) > 0.0

    def grant_caller_credit(self, api_key: str, account_id: str, amount: float) -> bool:
        """Grant the one-time Caller Credit for an account (US-039).

        The event id is derived from the account, not the key, so the grant is
        idempotent per account — across retries, additional keys, and hive
        gossip replication alike. Returns True only when credit was applied.
        """
        if amount <= 0:
            return False
        event_id = f"caller-credit-{account_id}"
        with self._lock:
            if event_id in self._seen_event_ids:
                return False
            self._apply_locked({
                "id": event_id,
                "type": "credit",
                "api_key": api_key,
                "amount": amount,
                "ts": time.time(),
                "note": "caller-credit",
            })
            return True

    def credit_client(self, api_key: str, amount: float, *, note: str = "deposit") -> float:
        if amount <= 0:
            raise ValueError("credit amount must be positive")
        with self._lock:
            self._apply_locked({
                "id": f"credit-{uuid.uuid4().hex}",
                "type": "credit",
                "api_key": api_key,
                "amount": amount,
                "ts": time.time(),
                "note": note,
            })
            return self._client_balances[api_key]

    def charge_request(
        self,
        api_key: str,
        model: str,
        total_tokens: int,
        node_work: list[tuple[str | None, int]],
        *,
        input_tokens: int | None = None,
        output_tokens: int | None = None,
    ) -> dict:
        """Debit the client and split the fee 90/10.

        With ``input_tokens``/``output_tokens`` (US-045) the cost is
        ``input·in_rate + output·out_rate``; without them, legacy behavior —
        ``total_tokens`` at the blended rate. Replayed/gossiped events apply
        their recorded ``cost``, so old events are unaffected either way.

        ``node_work`` is ``[(wallet_address | None, work_units), ...]`` for the
        nodes that served the request. Work units of nodes without a wallet
        accrue to the protocol cut — there is nowhere to pay them out.
        The client balance may dip below zero on the final request; the next
        request is then rejected by ``has_funds`` (post-pay drift, standard
        metered-billing behavior).
        """
        if input_tokens is not None or output_tokens is not None:
            in_rate, out_rate = self.prices_for(model)
            in_tokens = max(0, input_tokens or 0)
            out_tokens = max(0, output_tokens or 0)
            cost = (in_tokens * in_rate + out_tokens * out_rate) / 1000.0
            total_tokens = in_tokens + out_tokens
        else:
            cost = self.price_for(model) * max(0, total_tokens) / 1000.0
        total_work = sum(max(0, w) for _, w in node_work)
        node_pool = cost * self._node_share
        shares: dict[str, float] = {}
        if total_work > 0:
            for wallet, work in node_work:
                if wallet and work > 0:
                    shares[wallet] = shares.get(wallet, 0.0) + node_pool * work / total_work
        protocol_amount = cost - sum(shares.values())
        with self._lock:
            if api_key not in self._client_balances and self._starting_credit > 0.0:
                self._apply_locked({
                    "id": f"credit-{uuid.uuid4().hex}",
                    "type": "credit",
                    "api_key": api_key,
                    "amount": self._starting_credit,
                    "ts": time.time(),
                    "note": "caller-credit",
                })
            event = {
                "id": f"charge-{uuid.uuid4().hex}",
                "type": "charge",
                "api_key": api_key,
                "model": model,
                "total_tokens": total_tokens,
                **({"input_tokens": max(0, input_tokens or 0),
                    "output_tokens": max(0, output_tokens or 0)}
                   if (input_tokens is not None or output_tokens is not None) else {}),
                "cost": cost,
                "shares": shares,
                "protocol_amount": protocol_amount,
                "ts": time.time(),
            }
            self._apply_locked(event)
            return event

    def bind_wallet(self, api_key: str, wallet: str, *, admin_override: bool = False) -> dict:
        """Bind a client wallet pubkey to an API key (US-032 deposits, C6).

        A wallet already bound to a *different* api_key is left untouched
        unless ``admin_override`` is set — enforced in ``_apply_locked`` so a
        conflicting gossip ``bind`` event can never silently clobber an
        existing binding either. Callers should check ``event["rejected"]``.
        """
        with self._lock:
            event = {
                "id": f"bind-{uuid.uuid4().hex}",
                "type": "bind",
                "api_key": api_key,
                "wallet": wallet,
                "admin_override": admin_override,
                "ts": time.time(),
            }
            self._apply_locked(event)
            return event

    def api_key_for_wallet(self, wallet: str) -> str | None:
        with self._lock:
            return self._wallet_bindings.get(wallet)

    def credit_deposit(self, api_key: str, amount: float, signature: str) -> dict | None:
        """Credit an on-chain deposit exactly once.

        The event id embeds the transaction signature, so replayed or
        re-observed transfers — locally or via hive gossip — are no-ops.
        Returns None when the signature was already credited.
        """
        if amount <= 0:
            raise ValueError("deposit amount must be positive")
        event_id = f"deposit-{signature}"
        with self._lock:
            if event_id in self._seen_event_ids:
                return None
            event = {
                "id": event_id,
                "type": "credit",
                "api_key": api_key,
                "amount": amount,
                "signature": signature,
                "ts": time.time(),
                "note": "onchain-deposit",
            }
            self._apply_locked(event)
            return event

    def has_event(self, event_id: str) -> bool:
        with self._lock:
            return event_id in self._seen_event_ids

    # ---- settlement (US-033) ----

    def payables(
        self,
        *,
        threshold: float,
        max_period: float,
        dust_floor: float,
        now: float | None = None,
        exclude: set[str] | None = None,
    ) -> list[tuple[str, float]]:
        """Wallets due a payout: pending ≥ threshold OR pending age ≥ max_period,
        never below the dust floor (mining-pool standard, ADR-0015)."""
        now = now if now is not None else time.time()
        exclude = exclude or set()
        due: list[tuple[str, float]] = []
        with self._lock:
            for wallet, pending in self._node_pending.items():
                if wallet in exclude or pending < max(dust_floor, 1e-9):
                    continue
                if pending >= threshold:
                    due.append((wallet, pending))
                    continue
                since = self._pending_since.get(wallet)
                if since is not None and now - since >= max_period:
                    due.append((wallet, pending))
        return due

    def confirm_settlement(self, settlement_id: str, signature: str) -> dict:
        """Record the on-chain transaction signature for a settlement batch."""
        with self._lock:
            event = {
                "id": f"settlement-{settlement_id}",
                "type": "settlement",
                "settlement_id": settlement_id,
                "signature": signature,
                "ts": time.time(),
            }
            self._apply_locked(event)
            return event

    def unconfirmed_settlements(self) -> dict[str, list[tuple[str, float]]]:
        """Settlement batches whose pending was debited but whose transaction
        has not been confirmed — resend these (idempotent by settlement id)."""
        with self._lock:
            return {
                sid: list(s["payouts"])
                for sid, s in self._settlements.items()
                if s["signature"] is None and s["payouts"]
            }

    def settlement_history(self) -> list[dict]:
        with self._lock:
            return [
                {
                    "settlement_id": sid,
                    "signature": s["signature"],
                    "payouts": [
                        {"wallet": w, "amount": a} for w, a in s["payouts"]
                    ],
                    "ts": s["ts"],
                }
                for sid, s in sorted(self._settlements.items(), key=lambda kv: kv[1]["ts"])
            ]

    def settle_node_payout(self, wallet: str, amount: float, *, reference: str = "") -> dict:
        """Deduct a paid-out amount from a node's pending balance (US-033 hook).

        ADR-0015: clamps to whatever is still pending under the same lock as
        the deduction, so a forfeiture that lands between the settlement
        loop's `payables()` snapshot and this call is never paid out on top
        of — the wallet gets at most what remains, never a stale amount.
        """
        if amount <= 0:
            raise ValueError("payout amount must be positive")
        with self._lock:
            amount = min(amount, self._node_pending.get(wallet, 0.0))
            event = {
                "id": f"payout-{uuid.uuid4().hex}",
                "type": "payout",
                "wallet": wallet,
                "amount": amount,
                "reference": reference,
                "ts": time.time(),
            }
            self._apply_locked(event)
            return event

    def forfeit_pending(self, wallet: str, *, reason: str = "fraud") -> dict:
        """Forfeit a node's entire pending balance to the protocol cut (US-034 hook)."""
        with self._lock:
            event = {
                "id": f"forfeit-{uuid.uuid4().hex}",
                "type": "forfeit",
                "wallet": wallet,
                "amount": self._node_pending.get(wallet, 0.0),
                "reason": reason,
                "ts": time.time(),
            }
            self._apply_locked(event)
            return event

    # ---- replication ----

    def events_since(self, index: int) -> tuple[list[dict], int]:
        """Return events after ``index`` in the local log and the new cursor."""
        with self._lock:
            return list(self._event_log[index:]), len(self._event_log)

    def apply_events(self, events: list[dict]) -> int:
        """Apply peer events not yet seen locally. Returns how many applied."""
        applied = 0
        with self._lock:
            for event in events:
                event_id = event.get("id")
                if not event_id or event_id in self._seen_event_ids:
                    continue
                self._apply_locked(event)
                applied += 1
        return applied

    def _apply_locked(self, event: dict) -> None:
        etype = event.get("type")
        if etype == "credit":
            key = event["api_key"]
            self._client_balances[key] = self._client_balances.get(key, 0.0) + float(event["amount"])
        elif etype == "charge":
            key = event["api_key"]
            self._client_balances[key] = self._client_balances.get(key, 0.0) - float(event["cost"])
            for wallet, amount in (event.get("shares") or {}).items():
                self._node_pending[wallet] = self._node_pending.get(wallet, 0.0) + float(amount)
                if amount > 0:
                    self._pending_since.setdefault(wallet, float(event.get("ts", time.time())))
            self._protocol_cut += float(event.get("protocol_amount", 0.0))
        elif etype == "payout":
            wallet = event["wallet"]
            amount = float(event["amount"])
            ts = float(event.get("ts", time.time()))
            self._node_pending[wallet] = self._node_pending.get(wallet, 0.0) - amount
            self._pending_since.pop(wallet, None)
            self._last_payout_ts[wallet] = ts
            reference = event.get("reference") or ""
            if reference:
                settlement = self._settlements.setdefault(
                    reference, {"payouts": [], "signature": None, "ts": ts}
                )
                settlement["payouts"].append((wallet, amount))
        elif etype == "settlement":
            settlement = self._settlements.setdefault(
                event["settlement_id"],
                {"payouts": [], "signature": None, "ts": float(event.get("ts", 0.0))},
            )
            settlement["signature"] = event.get("signature")
        elif etype == "forfeit":
            wallet = event["wallet"]
            amount = float(event.get("amount", 0.0))
            self._node_pending[wallet] = self._node_pending.get(wallet, 0.0) - amount
            self._protocol_cut += amount
        elif etype == "bind":
            wallet = event["wallet"]
            existing = self._wallet_bindings.get(wallet)
            if existing is not None and existing != event["api_key"] and not event.get("admin_override"):
                event["rejected"] = True
            else:
                self._wallet_bindings[wallet] = event["api_key"]
        else:
            return
        self._seen_event_ids.add(event["id"])
        self._event_log.append(event)
        self._dirty = True

    # ---- views ----

    def get_client_balance(self, api_key: str) -> float:
        with self._lock:
            return self._client_balances.get(api_key, 0.0)

    def get_node_pending(self, wallet: str) -> float:
        with self._lock:
            return self._node_pending.get(wallet, 0.0)

    def usage_for(self, api_keys: list[str], *, recent_limit: int | None = None) -> dict:
        """Aggregate charge history for a set of API keys (dashboard view)."""
        keys = set(api_keys)
        requests = 0
        total_tokens = 0
        total_cost = 0.0
        records: list[dict] = []
        with self._lock:
            for event in self._event_log:
                if event.get("type") != "charge" or event.get("api_key") not in keys:
                    continue
                requests += 1
                total_tokens += int(event.get("total_tokens", 0))
                total_cost += float(event.get("cost", 0.0))
                records.append({
                    "api_key": event["api_key"],
                    "model": event.get("model"),
                    "total_tokens": event.get("total_tokens", 0),
                    "cost": event.get("cost", 0.0),
                    "ts": event.get("ts", 0.0),
                })
        recent = records[-recent_limit:] if recent_limit is not None else records
        return {
            "requests": requests,
            "total_tokens": total_tokens,
            "total_cost": total_cost,
            "records": records,
            "recent": recent,
        }

    def snapshot(self) -> dict:
        with self._lock:
            return {
                "clients": dict(self._client_balances),
                "node_pending": dict(self._node_pending),
                "wallet_bindings": dict(self._wallet_bindings),
                "protocol_cut": self._protocol_cut,
                "prices": dict(self._prices),
                "default_price_per_1k_tokens": self._default_price_per_1k,
                "node_share": self._node_share,
                "event_count": len(self._event_log),
                "forfeits": [
                    {
                        "wallet": e["wallet"],
                        "amount": e.get("amount", 0.0),
                        "reason": e.get("reason", ""),
                        "ts": e.get("ts", 0.0),
                    }
                    for e in self._event_log if e.get("type") == "forfeit"
                ][-20:],
            }

    # ---- persistence ----

    def _init_db(self) -> None:
        con = sqlite3.connect(self._db_path)  # type: ignore[arg-type]
        con.execute(
            "CREATE TABLE IF NOT EXISTS billing_events "
            "(event_id TEXT PRIMARY KEY, payload TEXT NOT NULL, ts REAL NOT NULL)"
        )
        con.commit()
        con.close()

    def _load_from_db(self) -> None:
        con = sqlite3.connect(self._db_path)  # type: ignore[arg-type]
        rows = con.execute(
            "SELECT payload FROM billing_events ORDER BY ts, event_id"
        ).fetchall()
        con.close()
        with self._lock:
            for (payload,) in rows:
                try:
                    event = json.loads(payload)
                except json.JSONDecodeError:
                    continue
                if event.get("id") not in self._seen_event_ids:
                    self._apply_locked(event)
            self._dirty = False

    def save_to_db(self) -> None:
        if not self._db_path:
            return
        with self._lock:
            if not self._dirty:
                return
            events = list(self._event_log)
            self._dirty = False
        con = sqlite3.connect(self._db_path)  # type: ignore[arg-type]
        con.executemany(
            "INSERT OR IGNORE INTO billing_events (event_id, payload, ts) VALUES (?, ?, ?)",
            [(e["id"], json.dumps(e), float(e.get("ts", 0.0))) for e in events],
        )
        con.commit()
        con.close()