from __future__ import annotations

from dataclasses import dataclass

from solders.keypair import Keypair

import wzrd.loop as loop_module
from wzrd.client import PickResult, WZRDError


@dataclass
class _FakeSession:
    token: str = "session-token"


class _FakeAgent:
    def __init__(self) -> None:
        self.auth_calls = 0
        self.claim_calls = 0
        self.report_calls = 0
        self.stake_calls = 0
        self.stake_amounts: list[int] = []
        self.stake_locks: list[int] = []
        self.status_responses: list[dict] = [
            {"solana": {"sol_lamports": 1_000_000_000, "ccm_balance": 0}},
        ]
        self.claim_status_response = {"claimable_ccm": 0, "relay_ready": True}
        self.claim_response = {"status": "already_claimed"}
        self.raise_on_status_at: set[int] = set()
        self.status_calls = 0

    def authenticate(self, keypair):
        self.auth_calls += 1
        return _FakeSession()

    def earned(self):
        return {
            "pending_ccm": 0,
            "total_earned_ccm": 0,
            "pipeline": {},
        }

    def status(self):
        self.status_calls += 1
        if self.status_calls in self.raise_on_status_at:
            raise WZRDError("status unavailable")
        if self.status_responses:
            return self.status_responses.pop(0)
        return {"solana": {"sol_lamports": 1_000_000_000, "ccm_balance": 0}}

    def infer(self, model, *, task_type=None, pick_id=None):
        return {}

    def report_pick(self, choice, **kwargs):
        self.report_calls += 1
        return {"pending_ccm": 0}

    def claim_status(self):
        return self.claim_status_response

    def claim(self):
        self.claim_calls += 1
        return self.claim_response

    def stake(self, amount, lock_duration, keypair=None, channel_config=None):
        self.stake_calls += 1
        self.stake_amounts.append(amount)
        self.stake_locks.append(lock_duration)
        return {"tx_sig": "stake-tx"}

    def invalidate_session(self):
        return None

    def health_check(self):
        return True


def test_run_loop_stops_without_sleeping_after_max_cycles(monkeypatch):
    signer = Keypair()
    fake_agent = _FakeAgent()
    sleep_calls: list[float] = []

    monkeypatch.setattr("wzrd.loop.load_keypair", lambda keypair=None: signer)
    monkeypatch.setattr("wzrd.loop.WZRDAgent", lambda: fake_agent)
    monkeypatch.setattr("wzrd.loop._fetch_ccm_price", lambda: None)
    monkeypatch.setattr("wzrd.loop.shortlist", lambda task, limit=3: [])
    monkeypatch.setattr(
        "wzrd.loop.pick_details",
        lambda task: PickResult(
            model="moonshotai/Kimi-K2.5",
            signal_model="moonshotai/Kimi-K2.5",
            task=task,
            policy="strict",
            score=1.0,
            raw_score=1.0,
            trend="surging",
            confidence="high",
            action="candidate",
            platform="huggingface",
            reason="test signal",
        ),
    )
    monkeypatch.setattr("wzrd.loop.time.sleep", lambda seconds: sleep_calls.append(seconds))

    loop_module.run_loop(tasks=["code"], cycle_seconds=30, max_cycles=1, quiet=True)

    assert fake_agent.auth_calls == 1
    assert fake_agent.claim_calls == 0
    assert fake_agent.report_calls == 1
    assert sleep_calls == []


def test_run_loop_auto_stakes_observed_post_claim_balance_delta(monkeypatch):
    signer = Keypair()
    fake_agent = _FakeAgent()
    sleep_calls: list[float] = []

    fake_agent.claim_status_response = {"claimable_ccm": 1_000_000_000, "relay_ready": True}
    fake_agent.claim_response = {"tx_sig": "claim-tx", "status": "ok"}
    fake_agent.status_responses = [
        {"solana": {"sol_lamports": 1_000_000_000, "ccm_balance": 0}},
        {"solana": {"sol_lamports": 1_000_000_000, "ccm_balance": 0}},
        {"solana": {"sol_lamports": 1_000_000_000, "ccm_balance": 995_000_000}},
    ]

    monkeypatch.setattr("wzrd.loop.load_keypair", lambda keypair=None: signer)
    monkeypatch.setattr("wzrd.loop.WZRDAgent", lambda: fake_agent)
    monkeypatch.setattr("wzrd.loop._fetch_ccm_price", lambda: None)
    monkeypatch.setattr("wzrd.loop.shortlist", lambda task, limit=3: [])
    monkeypatch.setattr(
        "wzrd.loop.pick_details",
        lambda task: PickResult(
            model="moonshotai/Kimi-K2.5",
            signal_model="moonshotai/Kimi-K2.5",
            task=task,
            policy="strict",
            score=1.0,
            raw_score=1.0,
            trend="surging",
            confidence="high",
            action="candidate",
            platform="huggingface",
            reason="test signal",
        ),
    )
    monkeypatch.setattr("wzrd.loop.time.sleep", lambda seconds: sleep_calls.append(seconds))

    loop_module.run_loop(tasks=["code"], cycle_seconds=30, max_cycles=1, quiet=True, stake=True)

    assert fake_agent.claim_calls == 1
    assert fake_agent.stake_calls == 1
    assert fake_agent.stake_amounts == [995_000_000]
    assert fake_agent.stake_locks == [216_000 * 7]
    assert sleep_calls == []


def test_run_loop_auto_stakes_after_delayed_post_claim_balance_update(monkeypatch):
    signer = Keypair()
    fake_agent = _FakeAgent()

    fake_agent.claim_status_response = {"claimable_ccm": 1_000_000_000, "relay_ready": True}
    fake_agent.claim_response = {"tx_sig": "claim-tx", "status": "ok"}
    fake_agent.status_responses = [
        {"solana": {"sol_lamports": 1_000_000_000, "ccm_balance": 0}},
        {"solana": {"sol_lamports": 1_000_000_000, "ccm_balance": 0}},
        {"solana": {"sol_lamports": 1_000_000_000, "ccm_balance": 0}},
        {"solana": {"sol_lamports": 1_000_000_000, "ccm_balance": 995_000_000}},
    ]

    monkeypatch.setattr("wzrd.loop.load_keypair", lambda keypair=None: signer)
    monkeypatch.setattr("wzrd.loop.WZRDAgent", lambda: fake_agent)
    monkeypatch.setattr("wzrd.loop._fetch_ccm_price", lambda: None)
    monkeypatch.setattr("wzrd.loop.shortlist", lambda task, limit=3: [])
    monkeypatch.setattr(
        "wzrd.loop.pick_details",
        lambda task: PickResult(
            model="moonshotai/Kimi-K2.5",
            signal_model="moonshotai/Kimi-K2.5",
            task=task,
            policy="strict",
            score=1.0,
            raw_score=1.0,
            trend="surging",
            confidence="high",
            action="candidate",
            platform="huggingface",
            reason="test signal",
        ),
    )
    monkeypatch.setattr("wzrd.loop.time.sleep", lambda seconds: None)

    loop_module.run_loop(tasks=["code"], cycle_seconds=30, max_cycles=1, quiet=True, stake=True)

    assert fake_agent.claim_calls == 1
    assert fake_agent.stake_calls == 1
    assert fake_agent.stake_amounts == [995_000_000]


def test_run_loop_skips_auto_stake_when_spendable_balance_is_unavailable(monkeypatch):
    signer = Keypair()
    fake_agent = _FakeAgent()
    sleep_calls: list[float] = []

    fake_agent.claim_status_response = {"claimable_ccm": 1_000_000_000, "relay_ready": True}
    fake_agent.claim_response = {"tx_sig": "claim-tx", "status": "ok"}
    fake_agent.status_responses = [
        {"solana": {"sol_lamports": 1_000_000_000, "ccm_balance": 0}},
        {"solana": {"sol_lamports": 1_000_000_000, "ccm_balance": 0}},
    ]
    fake_agent.raise_on_status_at = {3}

    monkeypatch.setattr("wzrd.loop.load_keypair", lambda keypair=None: signer)
    monkeypatch.setattr("wzrd.loop.WZRDAgent", lambda: fake_agent)
    monkeypatch.setattr("wzrd.loop._fetch_ccm_price", lambda: None)
    monkeypatch.setattr("wzrd.loop.shortlist", lambda task, limit=3: [])
    monkeypatch.setattr(
        "wzrd.loop.pick_details",
        lambda task: PickResult(
            model="moonshotai/Kimi-K2.5",
            signal_model="moonshotai/Kimi-K2.5",
            task=task,
            policy="strict",
            score=1.0,
            raw_score=1.0,
            trend="surging",
            confidence="high",
            action="candidate",
            platform="huggingface",
            reason="test signal",
        ),
    )
    monkeypatch.setattr("wzrd.loop.time.sleep", lambda seconds: sleep_calls.append(seconds))

    loop_module.run_loop(tasks=["code"], cycle_seconds=30, max_cycles=1, quiet=True, stake=True)

    assert fake_agent.claim_calls == 1
    assert fake_agent.stake_calls == 0
    assert sleep_calls == []
