"""MailShift — tests for the NEW features: 2FA, brute-force lockout,
worker queue cap, scheduling, and sync-mode dedup."""
import imaplib
import time
import uuid
from datetime import datetime, timedelta, timezone

import pyotp
import pytest
import requests

from conftest import API, mig_payload

LOCAL_API = "http://localhost:8001/api"


def _wait_status(client, mid, targets=("completed", "failed"), timeout=120):
    end = time.time() + timeout
    data = {}
    while time.time() < end:
        r = client.get(f"{API}/migrations/{mid}", timeout=30)
        assert r.status_code == 200, r.text
        data = r.json()
        if data["status"] in targets:
            return data
        time.sleep(2)
    return data


def _purge_dest(user="user2", password="pass2"):
    """Empty the destination mailbox so dedup counting is deterministic."""
    M = imaplib.IMAP4("127.0.0.1", 143)
    try:
        M.login(user, password)
        typ, data = M.list()
        for line in data or []:
            if b"\\Noselect" in line:
                continue
            name = line.rsplit(b" ", 1)[-1].decode().strip('"')
            if M.select('"%s"' % name)[0] != "OK":
                continue
            typ, res = M.uid("SEARCH", None, "ALL")
            if typ == "OK" and res and res[0]:
                M.uid("STORE", b",".join(res[0].split()), "+FLAGS", "(\\Deleted)")
                M.expunge()
    finally:
        try:
            M.logout()
        except Exception:
            pass


def _dest_email_count(user="user2", password="pass2"):
    """Count messages across all selectable folders of the destination mailbox."""
    M = imaplib.IMAP4("127.0.0.1", 143)
    try:
        M.login(user, password)
        typ, data = M.list()
        total = 0
        for line in data or []:
            if b"\\Noselect" in line:
                continue
            name = line.rsplit(b" ", 1)[-1].decode().strip('"')
            typ, sel = M.select('"%s"' % name, readonly=True)
            if typ == "OK" and sel and sel[0]:
                total += int(sel[0])
        return total
    finally:
        try:
            M.logout()
        except Exception:
            pass


# ---------- 2FA (throwaway user, never admin) ----------
class TestTwoFactor:
    @pytest.fixture(scope="class")
    def throwaway(self):
        email = f"TEST_2fa_{uuid.uuid4().hex[:8]}@example.com"
        pwd = "secret123"
        r = requests.post(f"{API}/auth/register",
                          json={"email": email, "password": pwd, "name": "2FA QA"}, timeout=30)
        assert r.status_code == 200, r.text
        d = r.json()
        assert d["user"]["totp_enabled"] is False
        return {"email": email, "password": pwd, "token": d["token"]}

    def test_2fa_setup_enable_login_disable(self, throwaway):
        h = {"Authorization": f"Bearer {throwaway['token']}"}

        # setup
        s = requests.post(f"{API}/auth/2fa/setup", headers=h, timeout=30)
        assert s.status_code == 200, s.text
        sd = s.json()
        secret = sd["secret"]
        assert isinstance(secret, str) and len(secret) >= 16
        assert sd["otpauth_uri"].startswith("otpauth://totp/")
        assert "MailShift" in sd["otpauth_uri"]
        assert sd["qr"].startswith("data:image/png;base64,")
        # not enabled until verified
        assert requests.get(f"{API}/auth/me", headers=h, timeout=30).json()["totp_enabled"] is False

        # enable with wrong code
        bad = requests.post(f"{API}/auth/2fa/enable", json={"code": "000000"}, headers=h, timeout=30)
        assert bad.status_code == 400, bad.text

        # enable with valid code
        totp = pyotp.TOTP(secret)
        ok = requests.post(f"{API}/auth/2fa/enable", json={"code": totp.now()}, headers=h, timeout=30)
        assert ok.status_code == 200, ok.text
        assert ok.json()["totp_enabled"] is True
        assert requests.get(f"{API}/auth/me", headers=h, timeout=30).json()["totp_enabled"] is True

        # setup again is rejected while enabled
        again = requests.post(f"{API}/auth/2fa/setup", headers=h, timeout=30)
        assert again.status_code == 400

        # login without code -> twofa_required, no token
        l1 = requests.post(f"{API}/auth/login",
                           json={"email": throwaway["email"], "password": throwaway["password"]},
                           timeout=30)
        assert l1.status_code == 200, l1.text
        assert l1.json() == {"twofa_required": True}

        # login with wrong code -> 401
        l2 = requests.post(f"{API}/auth/login",
                           json={"email": throwaway["email"], "password": throwaway["password"],
                                 "code": "123456"}, timeout=30)
        assert l2.status_code == 401
        assert "2FA" in l2.json().get("detail", "")

        # login with valid code -> token
        l3 = requests.post(f"{API}/auth/login",
                           json={"email": throwaway["email"], "password": throwaway["password"],
                                 "code": totp.now()}, timeout=30)
        assert l3.status_code == 200, l3.text
        assert isinstance(l3.json()["token"], str)
        assert l3.json()["user"]["totp_enabled"] is True

        # disable with wrong password -> 401
        h2 = {"Authorization": f"Bearer {l3.json()['token']}"}
        bd = requests.post(f"{API}/auth/2fa/disable", json={"password": "nope"}, headers=h2, timeout=30)
        assert bd.status_code == 401

        # disable with correct password
        dd = requests.post(f"{API}/auth/2fa/disable",
                           json={"password": throwaway["password"]}, headers=h2, timeout=30)
        assert dd.status_code == 200, dd.text
        assert dd.json()["totp_enabled"] is False

        # plain login works again
        l4 = requests.post(f"{API}/auth/login",
                           json={"email": throwaway["email"], "password": throwaway["password"]},
                           timeout=30)
        assert l4.status_code == 200
        assert "token" in l4.json()

    def test_2fa_endpoints_require_auth(self):
        for path, body in [("/auth/2fa/setup", {}), ("/auth/2fa/enable", {"code": "1"}),
                           ("/auth/2fa/disable", {"password": "x"})]:
            r = requests.post(f"{API}{path}", json=body, timeout=30)
            assert r.status_code == 401, f"{path} -> {r.status_code}"


# ---------- Brute-force lockout ----------
class TestBruteForce:
    def test_lockout_on_sixth_attempt_localhost(self):
        """Identifier = client IP + email; use localhost so the IP is stable."""
        email = f"TEST_bf_{uuid.uuid4().hex[:8]}@example.com"
        codes = []
        for _ in range(6):
            r = requests.post(f"{LOCAL_API}/auth/login",
                              json={"email": email, "password": "wrong"}, timeout=30)
            codes.append(r.status_code)
        assert codes[:5] == [401] * 5, codes
        assert codes[5] == 429, codes
        last = requests.post(f"{LOCAL_API}/auth/login",
                             json={"email": email, "password": "wrong"}, timeout=30)
        assert "Too many failed attempts" in last.json().get("detail", "")

    def test_lockout_blocks_correct_password(self, test_credentials):
        """A locked identifier must reject even the right password (separate email used)."""
        email = f"TEST_bf2_{uuid.uuid4().hex[:8]}@example.com"
        reg = requests.post(f"{LOCAL_API}/auth/register",
                            json={"email": email, "password": "secret123", "name": "BF"}, timeout=30)
        assert reg.status_code == 200
        for _ in range(5):
            requests.post(f"{LOCAL_API}/auth/login", json={"email": email, "password": "bad"}, timeout=30)
        r = requests.post(f"{LOCAL_API}/auth/login",
                          json={"email": email, "password": "secret123"}, timeout=30)
        assert r.status_code == 429, r.text
        # other accounts unaffected
        ok = requests.post(f"{LOCAL_API}/auth/login", json=test_credentials, timeout=30)
        assert ok.status_code == 200, ok.text


# ---------- Worker queue ----------
class TestQueue:
    def test_queue_endpoint(self, client):
        r = client.get(f"{API}/queue", timeout=30)
        assert r.status_code == 200, r.text
        d = r.json()
        assert d["max_workers"] == 3, d
        assert isinstance(d["queued"], int) and d["queued"] >= 0
        assert isinstance(d["running"], int) and d["running"] >= 0

    def test_queue_requires_auth(self):
        assert requests.get(f"{API}/queue", timeout=30).status_code == 401

    def test_five_concurrent_migrations_all_complete(self, client, created_ids):
        ids = []
        for i in range(5):
            r = client.post(f"{API}/migrations", json=mig_payload(f"TEST_q_{i}"), timeout=60)
            assert r.status_code == 200, r.text
            ids.append(r.json()["id"])
        created_ids.extend(ids)
        for mid in ids:
            s = client.post(f"{API}/migrations/{mid}/start", timeout=30)
            assert s.status_code == 200, s.text
            assert s.json()["status"] == "queued"

        q = client.get(f"{API}/queue", timeout=30).json()
        assert q["running"] <= q["max_workers"], q

        for mid in ids:
            d = _wait_status(client, mid, timeout=180)
            assert d["status"] == "completed", (mid, d.get("status"), d.get("error"))
            assert d["migrated_emails"] == 8, d
            assert d["failed_emails"] == 0, d


# ---------- Scheduling ----------
class TestScheduling:
    def test_create_scheduled_migration_and_auto_start(self, client, created_ids):
        when = (datetime.now(timezone.utc) + timedelta(seconds=25)).isoformat().replace("+00:00", "Z")
        p = mig_payload("TEST_sched_auto")
        p["scheduled_at"] = when
        r = client.post(f"{API}/migrations", json=p, timeout=60)
        assert r.status_code == 200, r.text
        m = r.json()
        created_ids.append(m["id"])
        assert m["status"] == "scheduled", m
        assert m["scheduled_at"] is not None

        st = client.get(f"{API}/stats", timeout=30).json()
        assert st["scheduled"] >= 1, st
        assert st["active"] >= 1, st

        d = _wait_status(client, m["id"], timeout=140)
        assert d["status"] == "completed", d
        assert d["migrated_emails"] == 8, d

    def test_past_schedule_is_pending(self, client, created_ids):
        p = mig_payload("TEST_sched_past")
        p["scheduled_at"] = (datetime.now(timezone.utc) - timedelta(hours=1)).isoformat()
        r = client.post(f"{API}/migrations", json=p, timeout=60)
        assert r.status_code == 200, r.text
        m = r.json()
        created_ids.append(m["id"])
        assert m["status"] == "pending", m
        assert m["scheduled_at"] is None, m

    def test_invalid_schedule_string_is_ignored(self, client, created_ids):
        p = mig_payload("TEST_sched_bad")
        p["scheduled_at"] = "not-a-date"
        r = client.post(f"{API}/migrations", json=p, timeout=60)
        assert r.status_code == 200, r.text
        m = r.json()
        created_ids.append(m["id"])
        assert m["status"] == "pending", m

    def test_start_now_on_scheduled_migration(self, client, created_ids):
        p = mig_payload("TEST_sched_startnow")
        p["scheduled_at"] = (datetime.now(timezone.utc) + timedelta(hours=2)).isoformat()
        r = client.post(f"{API}/migrations", json=p, timeout=60)
        mid = r.json()["id"]
        created_ids.append(mid)
        assert r.json()["status"] == "scheduled"

        s = client.post(f"{API}/migrations/{mid}/start", timeout=30)
        assert s.status_code == 200, s.text
        d = _wait_status(client, mid, timeout=120)
        assert d["status"] == "completed", d
        assert d["migrated_emails"] == 8, d


# ---------- Sync mode (delta re-run, Message-ID dedup) ----------
class TestSyncMode:
    def test_rerun_completed_does_not_duplicate(self, client, created_ids):
        _purge_dest()
        assert _dest_email_count() == 0
        r = client.post(f"{API}/migrations", json=mig_payload("TEST_sync_mig"), timeout=60)
        mid = r.json()["id"]
        created_ids.append(mid)
        client.post(f"{API}/migrations/{mid}/start", timeout=30)
        d = _wait_status(client, mid, timeout=180)
        assert d["status"] == "completed", d
        count_after_first = _dest_email_count()
        assert count_after_first == 8, f"dest count after first run = {count_after_first}"

        # sync re-run
        s = client.post(f"{API}/migrations/{mid}/start", timeout=30)
        assert s.status_code == 200, s.text
        time.sleep(1)
        d2 = _wait_status(client, mid, timeout=180)
        assert d2["status"] == "completed", d2
        assert d2["migrated_emails"] == 8, d2

        logs = client.get(f"{API}/migrations/{mid}/logs", timeout=30).json()
        assert any("Sync run started" in x["message"] for x in logs), [x["message"] for x in logs][-5:]

        count_after_sync = _dest_email_count()
        assert count_after_sync == 8, f"destination duplicated mail: {count_after_sync}"

    def test_start_twice_immediately_rejected(self, client, created_ids):
        r = client.post(f"{API}/migrations", json=mig_payload("TEST_dupstart"), timeout=60)
        mid = r.json()["id"]
        created_ids.append(mid)
        s1 = client.post(f"{API}/migrations/{mid}/start", timeout=30)
        s2 = client.post(f"{API}/migrations/{mid}/start", timeout=30)
        assert s1.status_code == 200
        assert s2.status_code in (200, 400), s2.text
        _wait_status(client, mid, timeout=120)