from aix import ai, jobs
from aix.db import Jsonb
from aix.ids import new_id
from aix.main import _handlers, run_one
from aix.providers import set_provider
from aix.providers.base import ProviderError
from aix.providers.fake import FakeProvider

from .conftest import add_script
from .helpers import make_pdf
from .test_pipeline_e2e import EXTRACT


def _setup_one_page(conn, storage, world):
    sid = add_script(conn, storage, world, make_pdf(["text"]))
    jobs.enqueue(conn, "rasterize", {"script_id": sid, "batch_id": world["batch"]}, tenant_id=world["tenant"])
    assert run_one(conn, _handlers())
    return sid


def test_dedupe_key_blocks_second_live_job(conn):
    a = jobs.enqueue(conn, "noop", {}, tenant_id=None, dedupe_key="x")
    b = jobs.enqueue(conn, "noop", {}, tenant_id=None, dedupe_key="x")
    assert a and b is None
    assert run_one(conn, {"noop": lambda c, j: {}})
    assert jobs.enqueue(conn, "noop", {}, tenant_id=None, dedupe_key="x")   # allowed once the first is done


def test_schema_failure_is_retried_with_error_feedback(conn, storage, world, fake):
    _setup_one_page(conn, storage, world)
    fake.push("ExtractionPage", {"kind": "answer", "legibility": 7, "elements": []})   # invalid: answer w/o elements
    fake.push("ExtractionPage", EXTRACT[2])
    assert run_one(conn, _handlers())
    assert conn.execute("SELECT status FROM jobs WHERE type='extract_page'").fetchone()["status"] == "done"
    assert "failed validation" in fake.calls[-1]["parts"][-1].text


def test_provider_outage_retries_then_fails_with_flag(conn, storage, world, fake):
    sid = _setup_one_page(conn, storage, world)
    for _ in range(3):
        fake.push("ExtractionPage", ProviderError("529 overloaded", retryable=True))
    for attempt in range(3):
        assert run_one(conn, _handlers())
        conn.execute("UPDATE jobs SET run_after=now() WHERE status='queued'")      # skip backoff wait
    job = conn.execute("SELECT status, attempts FROM jobs WHERE type='extract_page'").fetchone()
    assert job["status"] == "failed" and job["attempts"] == 3
    flag = conn.execute("SELECT type, severity FROM flags WHERE script_id=%s", (sid,)).fetchone()
    assert flag["type"] == "processing_failed" and flag["severity"] == "block"
    assert conn.execute("SELECT status FROM scripts WHERE id=%s", (sid,)).fetchone()["status"] == "failed"


def test_failed_attempt_rolls_back_partial_writes(conn, storage, world, fake):
    _setup_one_page(conn, storage, world)
    fake.push("ExtractionPage", ProviderError("bad request", retryable=False))
    assert run_one(conn, _handlers())
    assert conn.execute("SELECT count(*) AS n FROM elements").fetchone()["n"] == 0
    assert conn.execute("SELECT status FROM jobs WHERE type='extract_page'").fetchone()["status"] == "failed"


def test_failover_after_repeated_provider_errors(conn, storage, world, fake):
    backup = FakeProvider()
    set_provider("backup", backup)
    import aix.providers as reg
    reg._instances["backup"] = backup
    conn.execute("INSERT INTO settings (id, key, scope, value_json) VALUES (%s,'ai.extract.fallback','platform',%s)",
                 (new_id(), Jsonb({"provider": "backup", "model": "backup-vision"})))
    for _ in range(3):
        ai._note_failure("fake")
    backup.push("ExtractionPage", EXTRACT[2])
    _setup_one_page(conn, storage, world)
    assert run_one(conn, _handlers())
    info = conn.execute("SELECT model_info_json FROM pages").fetchone()["model_info_json"]
    assert info["provider"] == "fake" and info["model"] == "backup-vision"  # FakeProvider reports its own name
    assert info["failover_from"]["model"] == "fake-vision"
    assert not fake.calls


def test_stale_running_job_is_reaped(conn):
    jid = jobs.enqueue(conn, "noop", {}, tenant_id=None)
    conn.execute("UPDATE jobs SET status='running', locked_at=now() - interval '10 minutes' WHERE id=%s", (jid,))
    assert jobs.reap_stale(conn) == 1
    assert conn.execute("SELECT status FROM jobs WHERE id=%s", (jid,)).fetchone()["status"] == "queued"


def test_tenant_isolation_on_payload_mismatch(conn, storage, world, fake):
    """A job carrying another tenant's id must not touch this tenant's script."""
    sid = add_script(conn, storage, world, make_pdf(["text"]))
    other = new_id()
    conn.execute("INSERT INTO tenants (id, biddaloy_school_id, name) VALUES (%s,'SCH2','Other')", (other,))
    jobs.enqueue(conn, "rasterize", {"script_id": sid, "batch_id": world["batch"]}, tenant_id=other)
    assert run_one(conn, _handlers())
    assert conn.execute("SELECT status FROM jobs").fetchone()["status"] == "failed"
    assert conn.execute("SELECT count(*) AS n FROM pages").fetchone()["n"] == 0
    assert conn.execute("SELECT count(*) AS n FROM flags").fetchone()["n"] == 0     # no cross-tenant flag