import os
import sys
from pathlib import Path
import math
import pytest
from fastapi.testclient import TestClient

# Ensure required env vars for app import/config
os.environ.setdefault("SECRET_KEY", "x" * 32)
os.environ.setdefault("DATABASE_URL", "sqlite:////tmp/delphi_test.sqlite")

# Ensure repository root on sys.path for direct test runs
ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from app.main import app  # noqa: E402
from app.auth.security import get_current_user  # noqa: E402
from app.database.base import SessionLocal  # noqa: E402
from app.models.pensions import NumberTable, LifeTable  # noqa: E402


@pytest.fixture(scope="module")
def client():
    # Override auth to bypass JWT for these tests
    class _User:
        def __init__(self):
            self.id = "test"
            self.username = "tester"
            self.is_admin = True
            self.is_active = True

    app.dependency_overrides[get_current_user] = lambda: _User()

    try:
        yield TestClient(app)
    finally:
        app.dependency_overrides.pop(get_current_user, None)


def _seed_monthly_na_series():
    db = SessionLocal()
    try:
        # Clear and seed NumberTable for months 0..4 for both sexes, race A
        for m in range(0, 5):
            db.query(NumberTable).filter(NumberTable.month == m).delete()
        # NA series: [100000, 90000, 80000, 70000, 60000]
        series = [100000.0, 90000.0, 80000.0, 70000.0, 60000.0]
        for idx, na in enumerate(series):
            row = NumberTable(
                month=idx,
                na_am=na,  # male (all races)
                na_af=na,  # female (all races)
                na_aa=na,  # all (all races)
            )
            db.add(row)

        # Provide an LE row in case fallback is used inadvertently
        db.query(LifeTable).filter(LifeTable.age == 65).delete()
        db.add(LifeTable(age=65, le_am=20.0, le_af=22.0, le_aa=21.0))
        db.commit()
    finally:
        db.close()


def _monthly_rate(pct: float) -> float:
    return pow(1.0 + pct / 100.0, 1.0 / 12.0) - 1.0


def test_single_life_no_discount_no_cola(client: TestClient):
    _seed_monthly_na_series()

    payload = {
        "monthly_benefit": 1000.0,
        "term_months": 5,
        "start_age": 65,
        "sex": "M",
        "race": "A",
        "discount_rate": 0.0,
        "cola_rate": 0.0,
    }
    resp = client.post("/api/pensions/valuation/single-life", json=payload)
    assert resp.status_code == 200, resp.text
    # Survival probs: [1.0, 0.9, 0.8, 0.7, 0.6]
    # PV = 1000 * sum = 4000
    assert math.isclose(resp.json()["pv"], 4000.0, rel_tol=1e-6, abs_tol=0.01)


def test_joint_survivor_no_discount_no_cola(client: TestClient):
    _seed_monthly_na_series()

    payload = {
        "monthly_benefit": 1000.0,
        "term_months": 5,
        "participant_age": 65,
        "participant_sex": "M",
        "participant_race": "A",
        "spouse_age": 63,
        "spouse_sex": "F",
        "spouse_race": "A",
        "survivor_percent": 50.0,
        "discount_rate": 0.0,
        "cola_rate": 0.0,
    }
    resp = client.post("/api/pensions/valuation/joint-survivor", json=payload)
    assert resp.status_code == 200, resp.text
    data = resp.json()
    # p = [1, .9, .8, .7, .6]
    # p_both = p^2 = [1, .81, .64, .49, .36]
    # payment_t = 1000*p_both + 1000*0.5*(p - p_both) = 500*p + 500*p_both
    expected = 500.0 * (1 + 0.9 + 0.8 + 0.7 + 0.6) + 500.0 * (1 + 0.81 + 0.64 + 0.49 + 0.36)
    assert math.isclose(data["pv_total"], expected, rel_tol=1e-6, abs_tol=0.01)
    # Components
    expected_both = 1000.0 * (1 + 0.81 + 0.64 + 0.49 + 0.36)
    expected_surv = expected - expected_both
    assert math.isclose(data["pv_participant_component"], expected_both, rel_tol=1e-6, abs_tol=0.01)
    assert math.isclose(data["pv_survivor_component"], expected_surv, rel_tol=1e-6, abs_tol=0.01)


def test_joint_survivor_last_survivor_basis(client: TestClient):
    _seed_monthly_na_series()

    payload = {
        "monthly_benefit": 1000.0,
        "term_months": 5,
        "participant_age": 65,
        "participant_sex": "M",
        "participant_race": "A",
        "spouse_age": 63,
        "spouse_sex": "F",
        "spouse_race": "A",
        "survivor_percent": 50.0,  # should be ignored in last_survivor for total
        "discount_rate": 0.0,
        "cola_rate": 0.0,
        "survivor_basis": "last_survivor",
    }
    resp = client.post("/api/pensions/valuation/joint-survivor", json=payload)
    assert resp.status_code == 200, resp.text
    data = resp.json()
    # p = [1, .9, .8, .7, .6]; p_both = p^2; p_either = p_part + p_sp - p_both = 2p - p^2
    p = [1.0, 0.9, 0.8, 0.7, 0.6]
    p_both = [x * x for x in p]
    p_either = [2 * x - x * x for x in p]
    expected_total = 1000.0 * sum(p_either)
    assert math.isclose(data["pv_total"], expected_total, rel_tol=1e-6, abs_tol=0.01)


def test_interpolation_between_number_rows(client: TestClient):
    # Seed months 0 and 2 only; expect linear interpolation for month 1
    db = SessionLocal()
    try:
        from app.models.pensions import NumberTable, LifeTable
        for m in range(0, 3):
            db.query(NumberTable).filter(NumberTable.month == m).delete()
        # only months 0 and 2
        db.add(NumberTable(month=0, na_am=100000.0, na_af=100000.0, na_aa=100000.0))
        db.add(NumberTable(month=2, na_am=80000.0, na_af=80000.0, na_aa=80000.0))
        # Ensure LE exists but should not be used due to interpolation
        db.query(LifeTable).filter(LifeTable.age == 65).delete()
        db.add(LifeTable(age=65, le_am=30.0, le_af=30.0, le_aa=30.0))
        db.commit()
    finally:
        db.close()

    payload = {
        "monthly_benefit": 1000.0,
        "term_months": 3,
        "start_age": 65,
        "sex": "A",
        "race": "A",
        "discount_rate": 0.0,
        "cola_rate": 0.0,
        "interpolation_method": "linear",
    }
    resp = client.post("/api/pensions/valuation/single-life", json=payload)
    assert resp.status_code == 200, resp.text
    # survival at t: 0->1.0, 1->0.9 (interpolated), 2->0.8
    expected = 1000.0 * (1.0 + 0.9 + 0.8)
    assert math.isclose(resp.json()["pv"], expected, rel_tol=1e-6, abs_tol=0.01)


def test_interpolation_step_method(client: TestClient):
    # Seed months 0 and 3 only; with step interpolation, months 1 and 2 carry month 0 value
    db = SessionLocal()
    try:
        from app.models.pensions import NumberTable
        for m in range(0, 4):
            db.query(NumberTable).filter(NumberTable.month == m).delete()
        db.add(NumberTable(month=0, na_aa=100000.0))
        db.add(NumberTable(month=3, na_aa=70000.0))
        db.commit()
    finally:
        db.close()

    payload = {
        "monthly_benefit": 1000.0,
        "term_months": 4,
        "start_age": 65,
        "sex": "A",
        "race": "A",
        "discount_rate": 0.0,
        "cola_rate": 0.0,
        "interpolation_method": "step",
    }
    resp = client.post("/api/pensions/valuation/single-life", json=payload)
    assert resp.status_code == 200, resp.text
    # Series becomes NA: [100k, 100k, 100k, 70k] -> p: [1.0,1.0,1.0,0.7]
    expected = 1000.0 * (1.0 + 1.0 + 1.0 + 0.7)
    assert math.isclose(resp.json()["pv"], expected, rel_tol=1e-6, abs_tol=0.01)


def test_max_age_truncation(client: TestClient):
    _seed_monthly_na_series()
    # start_age 65, max_age 66 -> at most 12 months
    payload = {
        "monthly_benefit": 1000.0,
        "term_months": 24,
        "start_age": 65,
        "sex": "A",
        "race": "A",
        "discount_rate": 0.0,
        "cola_rate": 0.0,
        "max_age": 66,
    }
    resp = client.post("/api/pensions/valuation/single-life", json=payload)
    assert resp.status_code == 200
    # With our seed p falls by 0.1 each month; but we only have 5 months seeded; extend to 12
    db = SessionLocal()
    try:
        from app.models.pensions import NumberTable
        for m in range(5, 12):
            db.query(NumberTable).filter(NumberTable.month == m).delete()
            db.add(NumberTable(month=m, na_aa=100000.0 - 10000.0 * m))
        db.commit()
    finally:
        db.close()
    # Recompute with extension
    resp = client.post("/api/pensions/valuation/single-life", json=payload)
    assert resp.status_code == 200
    # p: 1.0,0.9,0.8,0.7,0.6,0.5,0.4,0.3,0.2,0.1,0.0 (clamped by interpolation), 0.0 -> sum first 10 effectively
    expected = 1000.0 * (1.0 + 0.9 + 0.8 + 0.7 + 0.6 + 0.5 + 0.4 + 0.3 + 0.2 + 0.1)
    assert math.isclose(resp.json()["pv"], expected, rel_tol=1e-6, abs_tol=0.1)


def test_single_life_with_discount_and_cola(client: TestClient):
    _seed_monthly_na_series()
    # discount 12% annual, COLA 12% annual
    d_m = _monthly_rate(12.0)
    g_m = _monthly_rate(12.0)
    probs = [1.0, 0.9, 0.8, 0.7, 0.6]
    monthly = 1000.0

    # Expected PV: sum( monthly * p[t] * (1+g)^t / (1+i)^t )
    expected = 0.0
    g = 1.0
    d = 1.0
    for t, p in enumerate(probs):
        if t == 0:
            g = 1.0
            d = 1.0
        else:
            g *= (1.0 + g_m)
            d *= (1.0 + d_m)
        expected += monthly * p * g / d

    payload = {
        "monthly_benefit": monthly,
        "term_months": 5,
        "start_age": 65,
        "sex": "A",
        "race": "A",
        "discount_rate": 12.0,
        "cola_rate": 12.0,
        "cola_mode": "monthly",
    }
    resp = client.post("/api/pensions/valuation/single-life", json=payload)
    assert resp.status_code == 200, resp.text
    assert math.isclose(resp.json()["pv"], expected, rel_tol=1e-6, abs_tol=0.05)


def test_single_life_cola_annual_prorated_and_cap(client: TestClient):
    _seed_monthly_na_series()
    # Extend months 0..11 to allow a full year
    db = SessionLocal()
    try:
        from app.models.pensions import NumberTable
        for m in range(5, 12):
            db.query(NumberTable).filter(NumberTable.month == m).delete()
            db.add(NumberTable(month=m, na_aa=100000.0 - 10000.0 * m))
        db.commit()
    finally:
        db.close()

    # annual COLA 6% but cap at 4%; annual prorated mode
    # Payments monthly for 6 months, p[t]=1,.9,.8,.7,.6,.5
    # Growth at t uses 4% cap prorated: factors ~ 1.0, 1+0.04*(1/12), 1+0.04*(2/12), ...
    monthly = 1000.0
    probs = [1.0, 0.9, 0.8, 0.7, 0.6, 0.5]
    a = 0.04
    expected = 0.0
    for t, p in enumerate(probs):
        growth = (1.0 + a * (t / 12.0))
        expected += monthly * p * growth

    payload = {
        "monthly_benefit": monthly,
        "term_months": 6,
        "start_age": 65,
        "sex": "A",
        "race": "A",
        "discount_rate": 0.0,
        "cola_rate": 6.0,
        "cola_mode": "annual_prorated",
        "cola_cap_percent": 4.0,
    }
    resp = client.post("/api/pensions/valuation/single-life", json=payload)
    assert resp.status_code == 200, resp.text
    assert math.isclose(resp.json()["pv"], expected, rel_tol=1e-6, abs_tol=0.05)


def test_single_life_quarterly_with_deferral(client: TestClient):
    _seed_monthly_na_series()
    # Quarterly payments (3 months per payment), starting after 2-month deferral
    # Horizon 10 months, payments at t=2,5,8 (assuming 0-indexed months)
    # Survival probs p[t] = [1, .9, .8, .7, .6, .5, .4, .3, .2, .1] for first 10 implied by seed extension
    # But our seed has only 5 months; extend seed
    db = SessionLocal()
    try:
        # extend NumberTable to months 0..9 with linear decrement 100k - 10k*m
        from app.models.pensions import NumberTable
        for m in range(5, 10):
            db.query(NumberTable).filter(NumberTable.month == m).delete()
            db.add(NumberTable(month=m, na_aa=100000.0 - 10000.0 * m))
        db.commit()
    finally:
        db.close()

    payload = {
        "monthly_benefit": 1000.0,
        "term_months": 10,
        "start_age": 65,
        "sex": "A",
        "race": "A",
        "discount_rate": 0.0,
        "cola_rate": 0.0,
        "defer_months": 2,
        "payment_period_months": 3,
    }
    resp = client.post("/api/pensions/valuation/single-life", json=payload)
    assert resp.status_code == 200, resp.text
    # p = [1.0, .9, .8, .7, .6, .5, .4, .3, .2, .1]
    # Payments at t=2,5,8: amount = monthly * 3 * p[t]
    expected = 1000.0 * 3.0 * (0.8 + 0.5 + 0.2)
    assert math.isclose(resp.json()["pv"], expected, rel_tol=1e-6, abs_tol=0.01)


def test_fractional_deferral_prorated_first_payment(client: TestClient):
    _seed_monthly_na_series()
    # Monthly payments, defer 0.5 month: first payment at month 1 with 50% of base
    # With p = [1.0, 0.9, 0.8, 0.7, 0.6]
    monthly = 1000.0
    payload = {
        "monthly_benefit": monthly,
        "term_months": 5,
        "start_age": 65,
        "sex": "A",
        "race": "A",
        "discount_rate": 0.0,
        "cola_rate": 0.0,
        "defer_months": 0.5,
        "payment_period_months": 1,
    }
    resp = client.post("/api/pensions/valuation/single-life", json=payload)
    assert resp.status_code == 200, resp.text
    # First payment at t=1: 0.5*1000*0.9 + subsequent: 1000*0.8 + 1000*0.7 + 1000*0.6
    expected = 0.5 * monthly * 0.9 + monthly * (0.8 + 0.7 + 0.6)
    assert math.isclose(resp.json()["pv"], expected, rel_tol=1e-6, abs_tol=0.01)


def test_joint_survivor_participant_only_commencement(client: TestClient):
    _seed_monthly_na_series()
    payload = {
        "monthly_benefit": 1000.0,
        "term_months": 5,
        "participant_age": 65,
        "participant_sex": "M",
        "participant_race": "A",
        "spouse_age": 63,
        "spouse_sex": "F",
        "spouse_race": "A",
        "survivor_percent": 50.0,
        "discount_rate": 0.0,
        "cola_rate": 0.0,
        "survivor_commence_participant_only": True,
    }
    resp = client.post("/api/pensions/valuation/joint-survivor", json=payload)
    assert resp.status_code == 200, resp.text
    data = resp.json()
    # p = [1, .9, .8, .7, .6]; p_both = p^2; p_sp_only replaced by p_part for survivor component
    p = [1.0, 0.9, 0.8, 0.7, 0.6]
    p_both = [x * x for x in p]
    surv_component = sum(1000.0 * 0.5 * x for x in p)  # using participant survival
    both_component = sum(1000.0 * x for x in p_both)
    expected_total = 1000.0 + (both_component - 1000.0) + surv_component  # first period total includes guarantee at t=0 (from previous tests)
    # Given our service guarantees only if certain_months > 0, here it's 0, so no guarantee. Recompute expected total accordingly
    expected_total = both_component + surv_component
    assert math.isclose(data["pv_participant_component"], both_component, rel_tol=1e-6, abs_tol=0.01)
    assert math.isclose(data["pv_survivor_component"], surv_component, rel_tol=1e-6, abs_tol=0.01)
    assert math.isclose(data["pv_total"], expected_total, rel_tol=1e-6, abs_tol=0.01)


def test_certain_period_guarantee_then_mortality(client: TestClient):
    _seed_monthly_na_series()
    # Extend 0..9 again
    db = SessionLocal()
    try:
        from app.models.pensions import NumberTable
        for m in range(5, 10):
            db.query(NumberTable).filter(NumberTable.month == m).delete()
            db.add(NumberTable(month=m, na_aa=100000.0 - 10000.0 * m))
        db.commit()
    finally:
        db.close()

    # Monthly payments, no deferral, 3 months certain
    payload = {
        "monthly_benefit": 1000.0,
        "term_months": 6,
        "start_age": 65,
        "sex": "A",
        "race": "A",
        "discount_rate": 0.0,
        "cola_rate": 0.0,
        "defer_months": 0,
        "payment_period_months": 1,
        "certain_months": 3,
    }
    resp = client.post("/api/pensions/valuation/single-life", json=payload)
    assert resp.status_code == 200, resp.text
    # p = [1.0, .9, .8, .7, .6, .5]
    # guaranteed first 3: 1000, 1000, 1000; then mortality-weighted: 700, 600, 500
    expected = 1000.0 * 3 + 1000.0 * (0.7 + 0.6 + 0.5)
    assert math.isclose(resp.json()["pv"], expected, rel_tol=1e-6, abs_tol=0.01)

    # Joint-survivor: same guarantee logic applies to total stream
    payload_js = {
        "monthly_benefit": 1000.0,
        "term_months": 4,
        "participant_age": 65,
        "participant_sex": "M",
        "participant_race": "A",
        "spouse_age": 63,
        "spouse_sex": "F",
        "spouse_race": "A",
        "survivor_percent": 50.0,
        "discount_rate": 0.0,
        "cola_rate": 0.0,
        "defer_months": 0,
        "payment_period_months": 1,
        "certain_months": 2,
    }
    resp = client.post("/api/pensions/valuation/joint-survivor", json=payload_js)
    assert resp.status_code == 200, resp.text
    data = resp.json()
    # p = [1.0, .9, .8, .7]
    # total mortality stream (no guarantee): 1000*p_both + 500*(p - p_both)
    p = [1.0, 0.9, 0.8, 0.7]
    p_both = [x * x for x in p]
    mort = [1000.0 * pb + 500.0 * (x - pb) for x, pb in zip(p, p_both)]
    total_expected = 1000.0 + 1000.0 + mort[2] + mort[3]
    assert math.isclose(data["pv_total"], total_expected, rel_tol=1e-6, abs_tol=0.01)


def test_batch_single_life_mixed_success_failure(client: TestClient):
    _seed_monthly_na_series()

    valid_item = {
        "monthly_benefit": 1000.0,
        "term_months": 5,
        "start_age": 65,
        "sex": "M",
        "race": "A",
        "discount_rate": 0.0,
        "cola_rate": 0.0,
    }
    invalid_item = {
        "monthly_benefit": -100.0,  # invalid
        "term_months": 5,
        "start_age": 65,
        "sex": "M",
        "race": "A",
        "discount_rate": 0.0,
        "cola_rate": 0.0,
    }
    payload = {
        "items": [valid_item, invalid_item, valid_item]
    }
    resp = client.post("/api/pensions/valuation/batch-single-life", json=payload)
    assert resp.status_code == 200, resp.text
    data = resp.json()["results"]
    assert len(data) == 3
    assert data[0]["success"] is True
    assert "pv" in data[0]["result"]
    assert math.isclose(data[0]["result"]["pv"], 4000.0, rel_tol=1e-6)
    assert data[1]["success"] is False
    assert "monthly_benefit must be non-negative" in (data[1]["error"] or "")
    assert data[2]["success"] is True
    assert math.isclose(data[2]["result"]["pv"], 4000.0, rel_tol=1e-6)

def test_batch_joint_survivor_mixed_success_failure(client: TestClient):
    _seed_monthly_na_series()

    valid_item = {
        "monthly_benefit": 1000.0,
        "term_months": 5,
        "participant_age": 65,
        "participant_sex": "M",
        "participant_race": "A",
        "spouse_age": 63,
        "spouse_sex": "F",
        "spouse_race": "A",
        "survivor_percent": 50.0,
        "discount_rate": 0.0,
        "cola_rate": 0.0,
    }
    invalid_item = {
        "monthly_benefit": 1000.0,
        "term_months": 5,
        "participant_age": 65,
        "participant_sex": "M",
        "participant_race": "A",
        "spouse_age": 63,
        "spouse_sex": "F",
        "spouse_race": "A",
        "survivor_percent": 150.0,  # invalid >100
        "discount_rate": 0.0,
        "cola_rate": 0.0,
    }
    payload = {
        "items": [valid_item, invalid_item]
    }
    resp = client.post("/api/pensions/valuation/batch-joint-survivor", json=payload)
    assert resp.status_code == 200, resp.text
    data = resp.json()["results"]
    assert len(data) == 2
    assert data[0]["success"] is True
    assert "pv_total" in data[0]["result"]
    p = [1.0, 0.9, 0.8, 0.7, 0.6]
    p_both = [x * x for x in p]
    expected = 500.0 * sum(p) + 500.0 * sum(p_both)
    assert math.isclose(data[0]["result"]["pv_total"], expected, rel_tol=1e-6)
    assert data[1]["success"] is False
    assert "survivor_percent must be between 0 and 100" in (data[1]["error"] or "")