"""Tests for conviction voting governance (src/primitives/conviction.py)."""

import math
import numpy as np
import pytest

from src.primitives.conviction import (
    ConvictionParams,
    Proposal,
    Voter,
    ConvictionSystem,
    trigger_threshold,
    update_conviction,
    max_conviction,
    conviction_at_time,
    epochs_to_fraction,
    stake,
    unstake,
    tick,
    generate_conviction_curves,
    get_governance_metrics,
)


# ---------- Helpers ----------

def make_system(n_voters: int = 3, holdings_each: float = 10_000.0) -> ConvictionSystem:
    """Return a ConvictionSystem with a few voters."""
    params = ConvictionParams(alpha=0.9, beta=0.2, rho=0.0025, min_age=3)
    system = ConvictionSystem(
        params=params,
        total_supply=n_voters * holdings_each,
        total_funds=50_000.0,
    )
    for i in range(n_voters):
        system.voters[f"voter_{i}"] = Voter(
            id=f"voter_{i}",
            holdings=holdings_each,
            sentiment=0.7,
        )
    return system


def add_proposal(
    system: ConvictionSystem,
    prop_id: str = "p1",
    funds_requested: float = 0.05,
) -> Proposal:
    prop = Proposal(
        id=prop_id,
        title="Test Proposal",
        funds_requested=funds_requested,
    )
    system.proposals[prop_id] = prop
    return prop


# ---------- ConvictionParams ----------

class TestConvictionParams:
    def test_defaults(self):
        p = ConvictionParams()
        assert p.alpha == 0.9
        assert p.beta == 0.2
        assert p.rho == 0.0025
        assert p.min_age == 3

    def test_half_life_formula(self):
        p = ConvictionParams(alpha=0.5)
        # T_half = -ln(2)/ln(0.5) = 1.0
        assert p.half_life == pytest.approx(1.0)

    def test_half_life_alpha_09(self):
        p = ConvictionParams(alpha=0.9)
        expected = -np.log(2) / np.log(0.9)
        assert p.half_life == pytest.approx(expected)

    def test_half_life_alpha_zero_is_inf(self):
        p = ConvictionParams(alpha=0.0)
        assert p.half_life == float("inf")

    def test_from_half_life_roundtrip(self):
        target_hl = 10.0
        p = ConvictionParams.from_half_life(target_hl)
        assert p.half_life == pytest.approx(target_hl, rel=1e-6)

    def test_from_half_life_kwargs_forwarded(self):
        p = ConvictionParams.from_half_life(10.0, beta=0.3, min_age=5)
        assert p.beta == 0.3
        assert p.min_age == 5


# ---------- update_conviction ----------

class TestUpdateConviction:
    def test_formula_y_t_plus_1(self):
        alpha = 0.9
        y_t = 100.0
        x = 50.0
        result = update_conviction(y_t, x, alpha)
        assert result == pytest.approx(alpha * y_t + x)

    def test_zero_initial_conviction(self):
        result = update_conviction(0.0, 100.0, 0.9)
        assert result == pytest.approx(100.0)

    def test_zero_staked(self):
        result = update_conviction(200.0, 0.0, 0.9)
        assert result == pytest.approx(0.9 * 200.0)

    def test_convergence_behavior(self):
        # Repeatedly applying with same stake should converge to max_conviction
        x = 100.0
        alpha = 0.9
        y = 0.0
        for _ in range(1000):
            y = update_conviction(y, x, alpha)
        assert y == pytest.approx(max_conviction(x, alpha), rel=1e-4)


# ---------- max_conviction ----------

class TestMaxConviction:
    def test_formula(self):
        x = 100.0
        alpha = 0.9
        assert max_conviction(x, alpha) == pytest.approx(x / (1 - alpha))

    def test_alpha_one_is_inf(self):
        assert max_conviction(100.0, 1.0) == float("inf")

    def test_higher_alpha_higher_max(self):
        x = 100.0
        assert max_conviction(x, 0.95) > max_conviction(x, 0.5)


# ---------- conviction_at_time ----------

class TestConvictionAtTime:
    def test_epoch_zero_equals_initial(self):
        result = conviction_at_time(100.0, 0.9, 0, initial_conviction=50.0)
        assert result == pytest.approx(50.0)

    def test_epoch_one_matches_update_formula(self):
        x = 100.0
        alpha = 0.9
        y0 = 50.0
        result = conviction_at_time(x, alpha, 1, y0)
        expected = update_conviction(y0, x, alpha)
        assert result == pytest.approx(expected)

    def test_matches_iterative_application(self):
        x = 75.0
        alpha = 0.85
        y0 = 0.0
        n = 20
        # Iterate manually
        y = y0
        for _ in range(n):
            y = update_conviction(y, x, alpha)
        assert conviction_at_time(x, alpha, n, y0) == pytest.approx(y, rel=1e-9)

    def test_approaches_max_conviction(self):
        x = 100.0
        alpha = 0.9
        y_large = conviction_at_time(x, alpha, 500)
        assert y_large == pytest.approx(max_conviction(x, alpha), rel=1e-3)


# ---------- epochs_to_fraction ----------

class TestEpochsToFraction:
    def test_half_is_half_life(self):
        alpha = 0.9
        p = ConvictionParams(alpha=alpha)
        t_half = epochs_to_fraction(0.5, alpha)
        assert t_half == pytest.approx(p.half_life, rel=1e-6)

    def test_fraction_zero_is_inf(self):
        assert epochs_to_fraction(0.0, 0.9) == float("inf")

    def test_fraction_one_is_inf(self):
        assert epochs_to_fraction(1.0, 0.9) == float("inf")

    def test_higher_fraction_takes_longer(self):
        alpha = 0.9
        t_50 = epochs_to_fraction(0.5, alpha)
        t_90 = epochs_to_fraction(0.9, alpha)
        assert t_90 > t_50


# ---------- trigger_threshold ----------

class TestTriggerThreshold:
    def test_formula(self):
        p = ConvictionParams(alpha=0.9, beta=0.2, rho=0.0025)
        supply = 100_000.0
        share = 0.05
        expected = p.rho * supply / ((1 - p.alpha) * (p.beta - share) ** 2)
        assert trigger_threshold(share, supply, p) == pytest.approx(expected)

    def test_share_at_beta_returns_inf(self):
        p = ConvictionParams(beta=0.2)
        result = trigger_threshold(0.2, 100_000.0, p)
        assert result == float("inf")

    def test_share_above_beta_returns_inf(self):
        p = ConvictionParams(beta=0.2)
        result = trigger_threshold(0.25, 100_000.0, p)
        assert result == float("inf")

    def test_larger_supply_higher_threshold(self):
        p = ConvictionParams()
        t1 = trigger_threshold(0.05, 100_000.0, p)
        t2 = trigger_threshold(0.05, 200_000.0, p)
        assert t2 > t1

    def test_smaller_share_lower_threshold(self):
        p = ConvictionParams(beta=0.2)
        t_small = trigger_threshold(0.01, 100_000.0, p)
        t_large = trigger_threshold(0.1, 100_000.0, p)
        assert t_small < t_large


# ---------- Voter stake/unstake ----------

class TestVoterStaking:
    def test_stake_adds_to_voter_stakes(self):
        system = make_system()
        add_proposal(system, "p1", 0.05)
        stake(system, "voter_0", "p1", 500.0)
        assert system.voters["voter_0"].stakes.get("p1", 0.0) == pytest.approx(500.0)

    def test_stake_capped_at_holdings(self):
        system = make_system(holdings_each=1_000.0)
        add_proposal(system, "p1")
        stake(system, "voter_0", "p1", 999_999.0)
        assert system.voters["voter_0"].stakes.get("p1", 0.0) <= 1_000.0

    def test_stake_on_unknown_voter_is_noop(self):
        system = make_system()
        add_proposal(system, "p1")
        result = stake(system, "nobody", "p1", 100.0)
        assert result is system  # No crash, no change

    def test_stake_on_unknown_proposal_is_noop(self):
        system = make_system()
        result = stake(system, "voter_0", "nonexistent", 100.0)
        assert result is system

    def test_unstake_removes_tokens(self):
        system = make_system()
        add_proposal(system, "p1")
        stake(system, "voter_0", "p1", 500.0)
        unstake(system, "voter_0", "p1", 200.0)
        assert system.voters["voter_0"].stakes.get("p1", 0.0) == pytest.approx(300.0)

    def test_unstake_all_removes_key(self):
        system = make_system()
        add_proposal(system, "p1")
        stake(system, "voter_0", "p1", 500.0)
        unstake(system, "voter_0", "p1", 500.0)
        assert "p1" not in system.voters["voter_0"].stakes

    def test_unstake_capped_at_staked(self):
        system = make_system()
        add_proposal(system, "p1")
        stake(system, "voter_0", "p1", 200.0)
        unstake(system, "voter_0", "p1", 999_999.0)
        assert system.voters["voter_0"].stakes.get("p1", 0.0) == pytest.approx(0.0)

    def test_stake_accumulates_on_multiple_calls(self):
        system = make_system(holdings_each=10_000.0)
        add_proposal(system, "p1")
        stake(system, "voter_0", "p1", 200.0)
        stake(system, "voter_0", "p1", 300.0)
        assert system.voters["voter_0"].stakes.get("p1", 0.0) == pytest.approx(500.0)

    def test_cannot_stake_more_than_available_across_proposals(self):
        system = make_system(holdings_each=1_000.0)
        add_proposal(system, "p1")
        add_proposal(system, "p2")
        stake(system, "voter_0", "p1", 800.0)
        stake(system, "voter_0", "p2", 500.0)  # Only 200 left
        assert system.voters["voter_0"].stakes.get("p2", 0.0) == pytest.approx(200.0)


# ---------- tick ----------

class TestTick:
    def test_tick_increments_epoch(self):
        system = make_system()
        tick(system)
        assert system.epoch == 1

    def test_tick_updates_conviction(self):
        system = make_system()
        add_proposal(system, "p1")
        stake(system, "voter_0", "p1", 500.0)
        tick(system)
        assert system.voters["voter_0"].convictions.get("p1", 0.0) > 0.0

    def test_tick_increments_proposal_age(self):
        system = make_system()
        add_proposal(system, "p1")
        tick(system)
        assert system.proposals["p1"].age == 1

    def test_proposal_does_not_pass_before_min_age(self):
        system = make_system(n_voters=5, holdings_each=100_000.0)
        add_proposal(system, "p1", funds_requested=0.001)
        # Stake heavily
        for vid in system.voters:
            stake(system, vid, "p1", 90_000.0)
        # Tick fewer times than min_age (=3)
        tick(system)
        tick(system)
        assert system.proposals["p1"].status == "candidate"

    def test_proposal_passes_after_sufficient_conviction(self):
        """Create conditions where a proposal accumulates enough conviction."""
        params = ConvictionParams(alpha=0.9, beta=0.5, rho=0.0001, min_age=1)
        system = ConvictionSystem(
            params=params,
            total_supply=1_000_000.0,
            total_funds=100_000.0,
        )
        # One whale voter with enormous holdings
        system.voters["whale"] = Voter(id="whale", holdings=1_000_000.0)
        prop = Proposal(id="p1", title="Easy Pass", funds_requested=0.001)
        system.proposals["p1"] = prop

        # Stake all holdings
        stake(system, "whale", "p1", 1_000_000.0)

        # Tick many times to build conviction far above threshold
        for _ in range(100):
            if system.proposals["p1"].status == "passed":
                break
            tick(system)

        assert system.proposals["p1"].status == "passed"
        assert len(system.passed_proposals) == 1


# ---------- generate_conviction_curves ----------

class TestGenerateConvictionCurves:
    def test_returns_correct_keys(self):
        curves = generate_conviction_curves(100.0, 0.9, epochs=50)
        assert "time" in curves
        assert "charge" in curves
        assert "discharge" in curves
        assert "max" in curves

    def test_correct_shape(self):
        epochs = 30
        curves = generate_conviction_curves(100.0, 0.9, epochs=epochs)
        assert len(curves["time"]) == epochs
        assert len(curves["charge"]) == epochs
        assert len(curves["discharge"]) == epochs
        assert len(curves["max"]) == epochs

    def test_charge_starts_near_zero(self):
        curves = generate_conviction_curves(100.0, 0.9, epochs=50)
        # At t=0, charge = x*(1 - alpha^0)/(1-alpha) = x*0/(1-alpha) = 0
        assert curves["charge"][0] == pytest.approx(0.0)

    def test_charge_approaches_max(self):
        curves = generate_conviction_curves(100.0, 0.9, epochs=500)
        y_max = max_conviction(100.0, 0.9)
        assert curves["charge"][-1] == pytest.approx(y_max, rel=1e-3)

    def test_discharge_starts_at_max(self):
        curves = generate_conviction_curves(100.0, 0.9, epochs=50)
        y_max = max_conviction(100.0, 0.9)
        assert curves["discharge"][0] == pytest.approx(y_max)

    def test_discharge_decays(self):
        curves = generate_conviction_curves(100.0, 0.9, epochs=50)
        assert curves["discharge"][-1] < curves["discharge"][0]

    def test_max_array_constant(self):
        curves = generate_conviction_curves(100.0, 0.9, epochs=20)
        assert np.all(curves["max"] == curves["max"][0])


# ---------- get_governance_metrics ----------

class TestGetGovernanceMetrics:
    def test_returns_expected_fields(self):
        system = make_system()
        m = get_governance_metrics(system)
        assert "epoch" in m
        assert "total_supply" in m
        assert "total_staked" in m
        assert "staking_ratio" in m
        assert "active_proposals" in m
        assert "passed_proposals" in m
        assert "total_voters" in m
        assert "proposals" in m

    def test_voter_count_correct(self):
        system = make_system(n_voters=4)
        m = get_governance_metrics(system)
        assert m["total_voters"] == 4

    def test_staking_ratio_zero_when_nothing_staked(self):
        system = make_system()
        m = get_governance_metrics(system)
        assert m["staking_ratio"] == pytest.approx(0.0)

    def test_staking_ratio_updates_after_stake(self):
        system = make_system(n_voters=1, holdings_each=10_000.0)
        add_proposal(system, "p1")
        stake(system, "voter_0", "p1", 5_000.0)
        m = get_governance_metrics(system)
        assert m["staking_ratio"] == pytest.approx(5_000.0 / 10_000.0)

    def test_active_proposals_count(self):
        system = make_system()
        add_proposal(system, "p1")
        add_proposal(system, "p2")
        m = get_governance_metrics(system)
        assert m["active_proposals"] == 2

    def test_passed_proposals_count(self):
        system = make_system()
        dummy_proposal = Proposal(id="done", title="Done", status="passed")
        system.passed_proposals.append(dummy_proposal)
        m = get_governance_metrics(system)
        assert m["passed_proposals"] == 1


# ---------- Full governance lifecycle ----------

class TestFullGovernanceLifecycle:
    def test_lifecycle(self):
        """Create proposals, stake, tick until passage."""
        params = ConvictionParams(alpha=0.9, beta=0.5, rho=0.0001, min_age=2)
        system = ConvictionSystem(
            params=params,
            total_supply=500_000.0,
            total_funds=100_000.0,
        )
        # Create voters
        for i in range(5):
            system.voters[f"v{i}"] = Voter(id=f"v{i}", holdings=100_000.0)

        # Create a low-threshold proposal
        prop = Proposal(id="easy", title="Easy", funds_requested=0.001)
        system.proposals["easy"] = prop

        # All voters stake
        for i in range(5):
            stake(system, f"v{i}", "easy", 80_000.0)

        # Tick until passed or timeout
        passed = False
        for _ in range(200):
            tick(system)
            if system.proposals["easy"].status == "passed":
                passed = True
                break

        assert passed
        assert len(system.passed_proposals) == 1

        # Verify metrics reflect state
        m = get_governance_metrics(system)
        assert m["passed_proposals"] == 1
        assert m["epoch"] > 0