"""US-034: hardened proof-of-work — pending-balance forfeiture penalty.

The pending balance is the fraud collateral (ADR-0015): a confirmed divergence
forfeits it in full to the protocol cut alongside the strike; three strikes ban
the wallet, whose remaining pending is never paid out.
"""

import json
import urllib.error
import urllib.request

import pytest

from meshnet_contracts import LocalSolanaContracts
from meshnet_node.server import StubNodeServer
from meshnet_tracker.billing import BillingLedger
from meshnet_tracker.server import TrackerServer
from meshnet_validator import ToplocAuditConfig, ValidatorProcess
from meshnet_validator.audit import build_activation_proofs

MODEL = "stub-model"


def _record_event(contracts, session_id: str, output: str, wallet: str) -> None:
    contracts.validation.record_completed_inference(
        session_id=session_id,
        model=MODEL,
        messages=[{"role": "user", "content": "2+2"}],
        observed_output=output,
        route_nodes=[{"wallet_address": wallet, "shard_end": 31}],
    )


@pytest.fixture
def reference_node():
    node = StubNodeServer(shard_start=0, shard_end=31, is_last_shard=True)
    port = node.start()
    yield f"http://127.0.0.1:{port}"
    node.stop()


def _reference_output(reference_url: str) -> str:
    """What the reference stub answers for the probe messages."""
    req = urllib.request.Request(
        f"{reference_url}/v1/infer",
        data=json.dumps({"messages": [{"role": "user", "content": "2+2"}]}).encode(),
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    with urllib.request.urlopen(req) as r:
        return json.loads(r.read())["text"]


def test_divergence_forfeits_pending_and_strikes(reference_node):
    contracts = LocalSolanaContracts()
    contracts.registry.submit_stake("wallet-bad", 500)
    ledger = BillingLedger(starting_credit=10.0, default_price_per_1k=0.02)
    ledger.charge_request("client", MODEL, 1000, [("wallet-bad", 12)])
    pending_before = ledger.get_node_pending("wallet-bad")
    assert pending_before > 0

    validator = ValidatorProcess(
        contracts=contracts,
        billing=ledger,
        reference_node_url=reference_node,
        sample_rate=1.0,
        slash_amount=100,
        random_seed=7,
    )
    _record_event(contracts, "sess-1", "fraudulent nonsense", "wallet-bad")
    receipts = validator.validate_once()

    assert len(receipts) == 1
    assert ledger.get_node_pending("wallet-bad") == pytest.approx(0.0)
    snapshot = ledger.snapshot()
    # forfeited amount landed in the protocol cut on top of the original 10%
    assert snapshot["protocol_cut"] == pytest.approx(0.02 * 0.10 + pending_before)
    assert contracts.registry.get_wallet("wallet-bad").strike_count == 1


def test_matching_output_forfeits_nothing(reference_node):
    contracts = LocalSolanaContracts()
    contracts.registry.submit_stake("wallet-good", 500)
    ledger = BillingLedger(starting_credit=10.0, default_price_per_1k=0.02)
    ledger.charge_request("client", MODEL, 1000, [("wallet-good", 12)])
    pending_before = ledger.get_node_pending("wallet-good")

    validator = ValidatorProcess(
        contracts=contracts,
        billing=ledger,
        reference_node_url=reference_node,
        sample_rate=1.0,
        random_seed=7,
    )
    _record_event(contracts, "sess-1", _reference_output(reference_node), "wallet-good")
    receipts = validator.validate_once()

    assert receipts == []
    assert ledger.get_node_pending("wallet-good") == pytest.approx(pending_before)
    assert contracts.registry.get_wallet("wallet-good").strike_count == 0


def test_three_strikes_bans_and_bad_node_loses_everything(reference_node):
    """Deliberately-bad node: every job is fraudulent, checks always sample.

    Earn → caught → forfeit, three times over; the third strike bans the
    wallet, and the tracker rejects its registration.
    """
    contracts = LocalSolanaContracts()
    contracts.registry.submit_stake("wallet-bad", 500)
    ledger = BillingLedger(starting_credit=10.0, default_price_per_1k=0.02)
    validator = ValidatorProcess(
        contracts=contracts,
        billing=ledger,
        reference_node_url=reference_node,
        sample_rate=1.0,
        strike_threshold=3,
        random_seed=7,
    )

    for i in range(3):
        ledger.charge_request("client", MODEL, 1000, [("wallet-bad", 12)])
        _record_event(contracts, f"sess-{i}", f"fraud-{i}", "wallet-bad")
        validator.validate_once()

    wallet = contracts.registry.get_wallet("wallet-bad")
    assert wallet.strike_count == 3
    assert wallet.banned
    assert ledger.get_node_pending("wallet-bad") == pytest.approx(0.0)

    # banned wallet cannot register with the tracker
    tracker = TrackerServer(contracts=contracts)
    port = tracker.start()
    try:
        req = urllib.request.Request(
            f"http://127.0.0.1:{port}/v1/nodes/register",
            data=json.dumps({
                "endpoint": "http://127.0.0.1:9",
                "shard_start": 0,
                "shard_end": 31,
                "hardware_profile": {},
                "wallet_address": "wallet-bad",
                "score": 1.0,
            }).encode(),
            headers={"Content-Type": "application/json"},
            method="POST",
        )
        with pytest.raises(urllib.error.HTTPError) as exc_info:
            urllib.request.urlopen(req)
        assert exc_info.value.code == 403
    finally:
        tracker.stop()


def test_forfeit_endpoint_requires_auth_and_forfeits():
    contracts = LocalSolanaContracts()
    ledger = BillingLedger(starting_credit=10.0, default_price_per_1k=0.02)
    ledger.charge_request("client", MODEL, 1000, [("wallet-x", 12)])
    pending = ledger.get_node_pending("wallet-x")

    tracker = TrackerServer(
        contracts=contracts, billing=ledger, validator_service_token="test-svc-token",
    )
    port = tracker.start()
    url = f"http://127.0.0.1:{port}/v1/billing/forfeit"
    body = json.dumps({"wallet": "wallet-x", "reason": "fraud"}).encode()
    try:
        req = urllib.request.Request(
            url, data=body, headers={"Content-Type": "application/json"}, method="POST",
        )
        with pytest.raises(urllib.error.HTTPError) as exc_info:
            urllib.request.urlopen(req)
        assert exc_info.value.code == 401

        req = urllib.request.Request(
            url, data=body,
            headers={"Content-Type": "application/json", "Authorization": "Bearer test-svc-token"},
            method="POST",
        )
        with urllib.request.urlopen(req) as r:
            result = json.loads(r.read())
        assert result["forfeited"] == pytest.approx(pending)
        assert result["strike_count"] == 1
        assert result["banned"] is False
        assert ledger.get_node_pending("wallet-x") == pytest.approx(0.0)
    finally:
        tracker.stop()


def test_probation_earns_nothing_then_earning_begins():
    """First N jobs accrue no pending balance; job N+1 earns (issue 08)."""
    contracts = LocalSolanaContracts(probationary_job_count=2)
    ledger = BillingLedger(starting_credit=10.0, default_price_per_1k=0.02)

    tracker = TrackerServer(
        model_presets={
            MODEL: {
                "layers_start": 0,
                "layers_end": 11,
                "bytes_per_layer": {"bfloat16": 30 * 1024 * 1024},
            }
        },
        contracts=contracts,
        billing=ledger,
    )
    port = tracker.start()
    tracker_url = f"http://127.0.0.1:{port}"

    stub = StubNodeServer(
        shard_start=0, shard_end=11, is_last_shard=True, model=MODEL, tracker_mode=True,
    )
    stub_port = stub.start()
    reg = json.dumps({
        "endpoint": f"http://127.0.0.1:{stub_port}",
        "shard_start": 0,
        "shard_end": 11,
        "model": MODEL,
        "hardware_profile": {},
        "score": 1.0,
        "tracker_mode": True,
        "wallet_address": "wallet-new",
    }).encode()

    def _chat():
        req = urllib.request.Request(
            f"{tracker_url}/v1/chat/completions",
            data=json.dumps({
                "model": MODEL,
                "messages": [{"role": "user", "content": "hi"}],
            }).encode(),
            headers={
                "Content-Type": "application/json",
                "Authorization": "Bearer probation-client",
            },
            method="POST",
        )
        with urllib.request.urlopen(req) as r:
            r.read()

    try:
        req = urllib.request.Request(
            f"{tracker_url}/v1/nodes/register",
            data=reg,
            headers={"Content-Type": "application/json"},
            method="POST",
        )
        with urllib.request.urlopen(req) as r:
            r.read()

        _chat()  # job 1: probation
        _chat()  # job 2: probation
        assert ledger.get_node_pending("wallet-new") == pytest.approx(0.0)
        assert contracts.registry.probationary_jobs_remaining("wallet-new") == 0

        _chat()  # job 3: earning begins (tokens are 0 in the stub, so the
        # charge event exists but the amount is 0 — assert via event log)
        events, _ = ledger.events_since(0)
        charges = [e for e in events if e["type"] == "charge"]
        assert len(charges) == 3
        assert charges[0]["shares"] == {} and charges[1]["shares"] == {}
        assert "wallet-new" in charges[2]["shares"]
    finally:
        stub.stop()
        tracker.stop()


class _FakeToploc:
    """Fingerprint = the activations themselves; verify is exact equality
    (same contract as tests/test_hop_bisection.py's fake backend)."""

    def build_proofs_base64(self, activations, *, decode_batching_size, topk, skip_prefill):
        return {
            "activation_fingerprint": tuple(tuple(row) for row in activations),
            "decode_batching_size": decode_batching_size,
            "topk": topk,
            "skip_prefill": skip_prefill,
        }

    def verify_proofs_base64(self, activations, proofs, *, decode_batching_size, topk, skip_prefill):
        return proofs == {
            "activation_fingerprint": tuple(tuple(row) for row in activations),
            "decode_batching_size": decode_batching_size,
            "topk": topk,
            "skip_prefill": skip_prefill,
        }


class _HopReferenceValidator(ValidatorProcess):
    """Stands in for the reference node: returns canned ground-truth
    activations at each hop cut-point instead of making an HTTP call."""

    def __init__(self, *args, reference_activations_by_hop, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self._reference_activations_by_hop = reference_activations_by_hop

    def _run_teacher_forced_prefill_hops(self, **kwargs):
        return self._reference_activations_by_hop


def _record_two_hop_event(
    contracts, session_id: str, *, hop0_activations, hop1_activations, config, backend,
) -> None:
    """Two-hop route (wallet-hop0 upstream, wallet-hop1 downstream), each hop
    carrying its own on-demand TOPLOC commitment (AH-007 bisection format)."""
    hop0_claim = build_activation_proofs(hop0_activations, config=config, backend=backend)
    hop1_claim = build_activation_proofs(hop1_activations, config=config, backend=backend)
    token_ids = [101, 202, 303]
    contracts.validation.record_completed_inference(
        session_id=session_id,
        model=MODEL,
        messages=[{"role": "user", "content": "2+2"}],
        observed_output="two-hop pipeline output",
        route_nodes=[
            {
                "wallet_address": "wallet-hop0",
                "shard_start": 0,
                "shard_end": 15,
                "toploc_proof": hop0_claim.as_mapping(),
                "claimed_token_ids": token_ids,
            },
            {
                "wallet_address": "wallet-hop1",
                "shard_start": 16,
                "shard_end": 31,
                "toploc_proof": hop1_claim.as_mapping(),
                "claimed_token_ids": token_ids,
            },
        ],
    )


def test_60_request_stream_bans_intermittent_first_hop_cheater_not_last_hop():
    """Integration (AH-010): a 60-request stream through a two-hop route where
    the *first* hop (not the last) cheats on 3 of the jobs. TOPLOC bisection
    (issue 07) must blame wallet-hop0 specifically -- the old last-hop-only
    heuristic would have blamed the innocent wallet-hop1 instead. Each catch
    forfeits wallet-hop0's pending balance and strikes it in the same
    validation cycle; the third strike bans it within the 60-request stream,
    and the settlement loop's payables() then excludes only the banned wallet
    while the honest downstream hop keeps earning (ADR-0015/ADR-0018).
    """
    config = ToplocAuditConfig(topk=2, decode_batching_size=16)
    backend = _FakeToploc()
    reference_hop0 = [[1.0, 2.0], [3.0, 4.0]]
    reference_hop1 = [[5.0, 6.0], [7.0, 8.0]]
    corrupted_hop0 = [[9.0, 9.0], [9.0, 9.0]]
    cheat_at = {20, 40, 60}

    contracts = LocalSolanaContracts()
    contracts.registry.submit_stake("wallet-hop0", 500)
    contracts.registry.submit_stake("wallet-hop1", 500)
    ledger = BillingLedger(starting_credit=1000.0, default_price_per_1k=0.02)

    validator = _HopReferenceValidator(
        contracts=contracts,
        billing=ledger,
        reference_node_url="http://reference.invalid",
        sample_rate=1.0,
        strike_threshold=3,
        random_seed=7,
        toploc_config=config,
        toploc_backend=backend,
        reference_activations_by_hop=[reference_hop0, reference_hop1],
    )

    banned_at: int | None = None
    for i in range(1, 61):
        ledger.charge_request(
            "client", MODEL, 1000, [("wallet-hop0", 6), ("wallet-hop1", 6)],
        )
        _record_two_hop_event(
            contracts,
            f"sess-{i}",
            hop0_activations=corrupted_hop0 if i in cheat_at else reference_hop0,
            hop1_activations=reference_hop1,
            config=config,
            backend=backend,
        )
        validator.validate_once()
        if banned_at is None and contracts.registry.get_wallet("wallet-hop0").banned:
            banned_at = i

    assert banned_at is not None and banned_at <= 60
    cheater = contracts.registry.get_wallet("wallet-hop0")
    honest = contracts.registry.get_wallet("wallet-hop1")
    assert cheater.strike_count == 3
    assert cheater.banned
    assert not honest.banned
    # Every catch blamed the true (first, non-last) culprit hop -- never the
    # innocent downstream wallet.
    assert ledger.get_node_pending("wallet-hop0") == pytest.approx(0.0)
    assert ledger.get_node_pending("wallet-hop1") > 0.0

    # Settlement-loop wiring (server.py `_settlement_loop`): banned wallets
    # are excluded from payables even though their (forfeited-to-zero)
    # balance would fail the dust floor anyway -- the honest hop is still due.
    banned_wallets = {
        wallet for wallet in ("wallet-hop0", "wallet-hop1")
        if contracts.registry.get_wallet(wallet).banned
    }
    due = dict(ledger.payables(threshold=0.0, max_period=0.0, dust_floor=0.0, exclude=banned_wallets))
    assert "wallet-hop0" not in due
    assert due.get("wallet-hop1", 0.0) > 0.0