"""Sign-in security: authenticator-app codes (app/core/totp.py), two-factor
sign-in and recovery codes (app/crud/mfa.py), lockout, session length and
the tenant's policy (app/crud/sign_in_policy.py), and single sign-on over
OpenID Connect (app/crud/oidc.py) against a fake identity provider."""

import base64
import hashlib
import time
import uuid
from datetime import datetime, timezone
from urllib.parse import parse_qs, urlsplit

import httpx
import pytest
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from jose import jwk, jwt
from sqlalchemy import select, text

from app.core import totp
from app.core.auth import create_access_token, decode_access_token
from tests.conftest import AsyncSessionLocal
from app.crud import oidc
from app.models.identity import Role, Tenant, User, UserRole
from tests.conftest import TEST_ADMIN_PASSWORD

# ---- authenticator codes ---------------------------------------------------------------------------------


def test_totp_matches_the_rfc_6238_test_vectors():
    secret = base64.b32encode(b"12345678901234567890").decode()
    assert totp.code_at(secret, 59 // 30) == "287082"
    assert totp.code_at(secret, 1111111109 // 30) == "081804"
    assert totp.code_at(secret, 2000000000 // 30) == "279037"


def test_a_code_is_accepted_once_and_only_near_now():
    secret = totp.new_secret()
    now = 1_800_000_000
    step = totp.current_step(now)
    assert totp.verify(secret, totp.code_at(secret, step), now=now) == step
    assert totp.verify(secret, totp.code_at(secret, step - 1), now=now) == step - 1  # clock drift
    assert totp.verify(secret, totp.code_at(secret, step - 3), now=now) is None
    assert totp.verify(secret, totp.code_at(secret, step), last_step=step, now=now) is None  # used already
    assert totp.verify(secret, "12 34", now=now) is None
    uri = totp.provisioning_uri(secret, "a@b.co", "FlowFabric (Acme)")
    assert uri.startswith("otpauth://totp/FlowFabric%20%28Acme%29%3Aa%40b.co?secret=" + secret)
    assert totp.qr_svg(uri).startswith("<svg")


# ---- helpers ------------------------------------------------------------------------------------------------


def _now_code(secret: str) -> str:
    return totp.code_at(secret, totp.current_step())


async def _login(client, tenant, password=TEST_ADMIN_PASSWORD, email=None):
    return await client.post("/auth/login", json={"tenant_id": str(tenant["tenant_id"]), "email": email or tenant["admin_email"], "password": password})


async def _set_policy(tenant_id, **fields):
    async with AsyncSessionLocal() as db:
        async with db.begin():
            t = await db.get(Tenant, tenant_id)
            t.security_settings = {**(t.security_settings or {}), **fields}


async def _enable_mfa(client, auth_headers) -> tuple[str, list[str]]:
    setup = (await client.post("/auth/mfa/setup", headers=auth_headers)).json()
    resp = await client.post("/auth/mfa/confirm", headers=auth_headers, json={"code": _now_code(setup["secret"])})
    assert resp.status_code == 200, resp.text
    return setup["secret"], resp.json()["recovery_codes"]


async def _clear_last_step(user_id):
    """A test signs in twice inside one 30-second step — let the same code count again."""
    async with AsyncSessionLocal() as db:
        async with db.begin():
            await db.execute(text("UPDATE users SET mfa_last_step = mfa_last_step - 2 WHERE id = :id"), {"id": user_id})


# ---- two-factor sign-in ------------------------------------------------------------------------------------


async def test_setting_up_two_factor_then_signing_in_with_it(client, auth_headers, test_tenant):
    assert (await client.get("/auth/mfa", headers=auth_headers)).json() == {"enabled": False, "enabled_at": None, "recovery_codes_left": 0, "required": False}
    setup = (await client.post("/auth/mfa/setup", headers=auth_headers)).json()
    assert setup["otpauth_uri"].startswith("otpauth://totp/") and setup["qr_svg"].startswith("<svg")
    # Not on until a code checks out — a wrong one changes nothing.
    assert (await client.post("/auth/mfa/confirm", headers=auth_headers, json={"code": "000000"})).status_code == 400
    assert (await _login(client, test_tenant)).json()["access_token"]

    codes = (await client.post("/auth/mfa/confirm", headers=auth_headers, json={"code": _now_code(setup["secret"])})).json()["recovery_codes"]
    assert len(codes) == 10 and len(set(codes)) == 10
    assert (await client.get("/auth/mfa", headers=auth_headers)).json()["recovery_codes_left"] == 10

    first = (await _login(client, test_tenant)).json()
    assert first["verification_required"] == "totp" and first["access_token"] is None
    challenge = {"tenant_id": str(test_tenant["tenant_id"]), "challenge_token": first["challenge_token"]}
    # The challenge isn't a session.
    assert (await client.get("/auth/mfa", headers={"Authorization": f"Bearer {first['challenge_token']}"})).status_code == 401
    assert (await client.post("/auth/login/verify-mfa", json={**challenge, "code": "000000"})).status_code == 400

    await _clear_last_step(test_tenant["admin_id"])
    done = (await client.post("/auth/login/verify-mfa", json={**challenge, "code": _now_code(setup["secret"])})).json()
    assert done["access_token"] and decode_access_token(done["access_token"])["amr"] == ["pwd", "otp"]
    # The same code again is refused — it was used.
    again = (await _login(client, test_tenant)).json()
    resp = await client.post("/auth/login/verify-mfa", json={"tenant_id": str(test_tenant["tenant_id"]), "challenge_token": again["challenge_token"], "code": _now_code(setup["secret"])})
    assert resp.status_code == 400 and "already used" in resp.json()["detail"]

    # A recovery code works once.
    resp = await client.post("/auth/login/verify-mfa", json={"tenant_id": str(test_tenant["tenant_id"]), "challenge_token": again["challenge_token"], "code": codes[0].upper()})
    assert resp.status_code == 200 and resp.json()["notice"] == "You used a recovery code — 9 left."
    third = (await _login(client, test_tenant)).json()
    resp = await client.post("/auth/login/verify-mfa", json={"tenant_id": str(test_tenant["tenant_id"]), "challenge_token": third["challenge_token"], "code": codes[0]})
    assert resp.status_code == 400


async def test_turning_two_factor_off_needs_a_code_and_isnt_allowed_when_required(client, auth_headers, test_tenant):
    secret, codes = await _enable_mfa(client, auth_headers)
    assert (await client.post("/auth/mfa/disable", headers=auth_headers, json={"code": "000000"})).status_code == 400
    new = (await client.post("/auth/mfa/recovery-codes", headers=auth_headers, json={"code": codes[1]})).json()["recovery_codes"]
    assert len(new) == 10 and codes[2] not in new

    await _set_policy(test_tenant["tenant_id"], mfa_policy="internal")
    resp = await client.post("/auth/mfa/disable", headers=auth_headers, json={"code": new[0]})
    assert resp.status_code == 409
    await _set_policy(test_tenant["tenant_id"], mfa_policy="optional")
    resp = await client.post("/auth/mfa/disable", headers=auth_headers, json={"code": new[0]})
    assert resp.status_code == 200 and resp.json()["enabled"] is False
    assert (await _login(client, test_tenant)).json()["access_token"]


async def test_a_required_second_factor_is_set_up_during_sign_in(client, test_tenant):
    await _set_policy(test_tenant["tenant_id"], mfa_policy="internal")
    first = (await _login(client, test_tenant)).json()
    assert first["verification_required"] == "mfa_setup"
    challenge = {"tenant_id": str(test_tenant["tenant_id"]), "challenge_token": first["challenge_token"]}
    # A setup challenge can't be used to skip straight to a code.
    assert (await client.post("/auth/login/verify-mfa", json={**challenge, "code": "123456"})).status_code == 400

    setup = (await client.post("/auth/login/mfa-setup", json=challenge)).json()
    done = (await client.post("/auth/login/mfa-setup/confirm", json={**challenge, "code": _now_code(setup["secret"])})).json()
    assert done["access_token"] and len(done["recovery_codes"]) == 10
    assert (await _login(client, test_tenant)).json()["verification_required"] == "totp"


async def test_an_administrator_can_reset_someones_two_factor(client, auth_headers, test_tenant):
    await _enable_mfa(client, auth_headers)
    resp = await client.post(f"/users/{test_tenant['admin_id']}/reset-mfa", headers=auth_headers)
    assert resp.status_code == 200 and resp.json()["mfa_enabled_at"] is None
    assert (await client.post(f"/users/{test_tenant['admin_id']}/reset-mfa", headers=auth_headers)).status_code == 409
    assert (await _login(client, test_tenant)).json()["access_token"]


# ---- lockout and session length ----------------------------------------------------------------------------


async def test_repeated_failures_lock_the_account_until_it_passes_or_is_unlocked(client, auth_headers, test_tenant):
    await _set_policy(test_tenant["tenant_id"], lockout_attempts=3, lockout_minutes=10)
    for _ in range(2):
        assert (await _login(client, test_tenant, password="wrong")).status_code == 401
    resp = await _login(client, test_tenant, password="wrong")
    assert resp.status_code == 401 and "locked for 10 more minutes" in resp.json()["detail"]
    # Even the right password waits it out.
    resp = await _login(client, test_tenant)
    assert resp.status_code == 403 and resp.json()["detail"]["code"] == "locked"

    resp = await client.post(f"/users/{test_tenant['admin_id']}/unlock", headers=auth_headers)
    assert resp.status_code == 200 and resp.json()["locked_until"] is None
    assert (await _login(client, test_tenant)).json()["access_token"]


async def test_the_workspace_sets_how_long_a_session_lasts(client, test_tenant):
    await _set_policy(test_tenant["tenant_id"], session_minutes=30)
    token = (await _login(client, test_tenant)).json()["access_token"]
    lifetime = decode_access_token(token)["exp"] - time.time()
    assert 25 * 60 < lifetime <= 30 * 60


# ---- the policy's settings ------------------------------------------------------------------------------------


def _sso(**over):
    return {"enabled": True, "issuer": "https://idp.example.test", "client_id": "flowfabric", "client_secret": "s3cret-value", **over}


async def test_security_settings_keep_the_client_secret_to_themselves(client, auth_headers, test_tenant):
    got = (await client.get("/tenants/me/security", headers=auth_headers)).json()
    assert got["mfa_policy"] == "optional" and got["sso"]["enabled"] is False and got["sso_redirect_uri"].endswith("/auth/sso/callback")

    resp = await client.put("/tenants/me/security", headers=auth_headers, json={"sso": _sso(client_secret=None)})
    assert resp.status_code == 422  # no secret yet
    resp = await client.put("/tenants/me/security", headers=auth_headers, json={"mfa_policy": "internal", "sso": _sso()})
    assert resp.status_code == 200, resp.text
    body = resp.json()
    assert body["sso"]["has_client_secret"] is True and "s3cret" not in resp.text
    assert "s3cret" not in (await client.get("/tenants/me", headers=auth_headers)).text
    async with AsyncSessionLocal() as db:
        stored = (await db.get(Tenant, test_tenant["tenant_id"])).security_settings
    assert "s3cret" not in str(stored) and stored["sso"]["client_secret_encrypted"]
    # Saving again without a secret keeps the one stored.
    resp = await client.put("/tenants/me/security", headers=auth_headers, json={"mfa_policy": "internal", "sso": _sso(client_secret=None, button_label="Sign in with Acme")})
    assert resp.status_code == 200 and resp.json()["sso"]["has_client_secret"] is True

    # Requiring SSO needs a session that signed in by SSO.
    resp = await client.put("/tenants/me/security", headers=auth_headers, json={"sso": _sso(require_sso=True)})
    assert resp.status_code == 409
    sso_token = create_access_token(tenant_id=test_tenant["tenant_id"], user_id=test_tenant["admin_id"], vendor_id=None, amr=["sso"])
    resp = await client.put("/tenants/me/security", headers={"Authorization": f"Bearer {sso_token}"}, json={"sso": _sso(require_sso=True)})
    assert resp.status_code == 200
    resp = await _login(client, test_tenant)
    assert resp.status_code == 403 and resp.json()["detail"]["code"] == "sso_required"
    assert (await client.get(f"/auth/sso-config?tenant_id={test_tenant['tenant_id']}")).json() == {"enabled": True, "button_label": "Sign in with SSO", "required": True, "password_reset": False}


# ---- single sign-on against a fake identity provider -----------------------------------------------------------

ISSUER = "https://idp.example.test"
RETURN_TO = "http://localhost:5173/sso/callback"


class FakeIdP:
    """Discovery, keys and a token endpoint, as a real provider serves them."""

    def __init__(self):
        self.key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
        self.pem = self.key.private_bytes(serialization.Encoding.PEM, serialization.PrivateFormat.PKCS8, serialization.NoEncryption())
        public = self.key.public_key().public_bytes(serialization.Encoding.PEM, serialization.PublicFormat.SubjectPublicKeyInfo)
        self.jwk = {**jwk.construct(public, "RS256").to_dict(), "kid": "k1", "use": "sig"}
        self.claims: dict = {}
        self.codes: dict[str, dict] = {}  # code -> what the authorize request asked for
        self.sign_with = self.pem

    def authorize(self, url: str, code: str = "good-code") -> dict:
        q = {k: v[0] for k, v in parse_qs(urlsplit(url).query).items()}
        self.codes[code] = q
        return q

    def id_token(self, request: dict, **over) -> str:
        now = int(time.time())
        claims = {"iss": ISSUER, "aud": "flowfabric", "sub": "idp-user-1", "email": "", "email_verified": True, "nonce": request["nonce"], "iat": now, "exp": now + 300, **self.claims, **over}
        return jwt.encode(claims, self.sign_with, algorithm="RS256", headers={"kid": "k1"})

    def handler(self, request: httpx.Request) -> httpx.Response:
        path = request.url.path
        if path == "/.well-known/openid-configuration":
            return httpx.Response(200, json={"issuer": ISSUER, "authorization_endpoint": f"{ISSUER}/authorize", "token_endpoint": f"{ISSUER}/token", "jwks_uri": f"{ISSUER}/jwks"})
        if path == "/jwks":
            return httpx.Response(200, json={"keys": [self.jwk]})
        if path == "/token":
            form = {k: v[0] for k, v in parse_qs(request.content.decode()).items()}
            asked = self.codes.pop(form.get("code"), None)
            challenge = base64.urlsafe_b64encode(hashlib.sha256(form.get("code_verifier", "").encode()).digest()).decode().rstrip("=")
            if asked is None or form.get("client_secret") != "s3cret-value" or challenge != asked["code_challenge"] or form["redirect_uri"] != asked["redirect_uri"]:
                return httpx.Response(400, json={"error": "invalid_grant", "error_description": "bad code"})
            return httpx.Response(200, json={"id_token": self.id_token(asked), "token_type": "Bearer"})
        return httpx.Response(404)


@pytest.fixture
async def idp(monkeypatch, client, auth_headers, test_tenant):
    fake = FakeIdP()
    fake.claims = {"email": test_tenant["admin_email"]}
    monkeypatch.setattr(oidc, "_http", lambda: httpx.AsyncClient(transport=httpx.MockTransport(fake.handler)))
    monkeypatch.setattr(oidc, "_check_address", lambda url: None)  # the fake host doesn't resolve
    oidc._discovery.clear()
    oidc._jwks.clear()
    resp = await client.put("/tenants/me/security", headers=auth_headers, json={"sso": _sso()})
    assert resp.status_code == 200, resp.text
    return fake


async def _sso_round_trip(client, test_tenant, idp, code="good-code"):
    start = await client.get("/auth/sso/start", params={"tenant_id": str(test_tenant["tenant_id"]), "return_to": RETURN_TO})
    assert start.status_code == 302, start.text
    assert start.headers["location"].startswith(f"{ISSUER}/authorize?")
    asked = idp.authorize(start.headers["location"], code)
    assert asked["code_challenge_method"] == "S256" and asked["redirect_uri"].endswith("/auth/sso/callback")
    return asked, await client.get("/auth/sso/callback", params={"state": asked["state"], "code": code})


def _token_from(resp) -> str:
    location = resp.headers["location"]
    assert location.startswith(RETURN_TO + "#sso_token="), location
    return location.split("#sso_token=", 1)[1]


def _error_from(resp) -> str:
    q = parse_qs(urlsplit(resp.headers["location"]).query)
    return q["sso_error"][0]


async def test_single_sign_on_signs_in_and_links_the_account(client, test_tenant, idp):
    asked, resp = await _sso_round_trip(client, test_tenant, idp)
    token = _token_from(resp)
    claims = decode_access_token(token)
    assert claims["amr"] == ["sso"] and claims["user_id"] == str(test_tenant["admin_id"])
    session = (await client.get("/auth/session", headers={"Authorization": f"Bearer {token}"})).json()
    assert session["user"]["email"] == test_tenant["admin_email"] and session["tenant"]["id"] == str(test_tenant["tenant_id"])
    async with AsyncSessionLocal() as db:
        assert (await db.get(User, test_tenant["admin_id"])).external_auth_id == "oidc:idp-user-1"

    # The attempt was used up.
    replay = await client.get("/auth/sso/callback", params={"state": asked["state"], "code": "good-code"})
    assert replay.status_code == 400 and "already used" in replay.text


async def test_single_sign_on_refuses_what_doesnt_check_out(client, test_tenant, idp):
    # Signed by a key the provider doesn't publish.
    other = rsa.generate_private_key(public_exponent=65537, key_size=2048)
    idp.sign_with = other.private_bytes(serialization.Encoding.PEM, serialization.PrivateFormat.PKCS8, serialization.NoEncryption())
    _, resp = await _sso_round_trip(client, test_tenant, idp)
    assert "didn't check out" in _error_from(resp)
    idp.sign_with = idp.pem

    for claims, expected in (
        ({"aud": "someone-else"}, "didn't check out"),
        ({"iss": "https://evil.example"}, "didn't check out"),
        ({"exp": int(time.time()) - 60}, "didn't check out"),
        ({"nonce": "not-this-attempt"}, "doesn't belong to this sign-in attempt"),
        ({"email_verified": False}, "hasn't verified"),
        ({"email": "stranger@elsewhere.test", "sub": "idp-user-2"}, "no account for stranger@elsewhere.test"),
    ):
        idp.claims = {"email": test_tenant["admin_email"], **claims}
        _, resp = await _sso_round_trip(client, test_tenant, idp)
        assert expected in _error_from(resp), (claims, _error_from(resp))

    # A wrong code (or verifier) at the token endpoint.
    start = await client.get("/auth/sso/start", params={"tenant_id": str(test_tenant["tenant_id"]), "return_to": RETURN_TO})
    asked = idp.authorize(start.headers["location"], "good-code")
    resp = await client.get("/auth/sso/callback", params={"state": asked["state"], "code": "forged-code"})
    assert "didn't accept" in _error_from(resp)
    # The provider reporting an error.
    start = await client.get("/auth/sso/start", params={"tenant_id": str(test_tenant["tenant_id"]), "return_to": RETURN_TO})
    asked = idp.authorize(start.headers["location"])
    resp = await client.get("/auth/sso/callback", params={"state": asked["state"], "error": "access_denied", "error_description": "User cancelled"})
    assert _error_from(resp) == "The identity provider said: User cancelled"
    # Only the app's own pages get a session.
    resp = await client.get("/auth/sso/start", params={"tenant_id": str(test_tenant["tenant_id"]), "return_to": "https://attacker.example/steal"})
    assert resp.status_code == 400


async def test_single_sign_on_can_create_accounts_for_allowed_domains(client, auth_headers, test_tenant, idp):
    async with AsyncSessionLocal() as db:
        role_id = (await db.execute(select(Role.id).where(Role.tenant_id == test_tenant["tenant_id"], Role.scope == "internal").limit(1))).scalar_one()
    resp = await client.put(
        "/tenants/me/security", headers=auth_headers,
        json={"sso": _sso(client_secret=None, auto_create_users=True, default_role_id=str(role_id), allowed_domains=["@Acme.test"])},
    )
    assert resp.status_code == 200 and resp.json()["sso"]["allowed_domains"] == ["acme.test"]

    idp.claims = {"email": "new.person@acme.test", "sub": "idp-new", "name": "New Person"}
    _, resp = await _sso_round_trip(client, test_tenant, idp)
    user_id = decode_access_token(_token_from(resp))["user_id"]
    async with AsyncSessionLocal() as db:
        user = await db.get(User, uuid.UUID(user_id))
        roles = (await db.execute(select(UserRole.role_id).where(UserRole.user_id == user.id))).scalars().all()
    assert (user.name, user.email, user.external_auth_id, user.password_hash) == ("New Person", "new.person@acme.test", "oidc:idp-new", None)
    assert roles == [role_id]

    idp.claims = {"email": "someone@other.test", "sub": "idp-x"}
    _, resp = await _sso_round_trip(client, test_tenant, idp)
    assert _error_from(resp) == "Accounts from other.test can't sign in to this workspace."