import json
import sqlite3
from unittest.mock import MagicMock

from grantflow.core.stores import (
    InMemoryJobStore,
    RedisJobStore,
    SQLiteIngestAuditStore,
    SQLiteJobStore,
    open_sqlite_connection,
)
from grantflow.core.strategies.factory import DonorFactory
from grantflow.swarm.hitl import HITLCheckpoint, HITLStatus


def test_inmemory_job_store_normalizes_state_contract_on_set_and_update():
    store = InMemoryJobStore()
    store.set(
        "job-inmem-1",
        {
            "status": "accepted",
            "state": {
                "donor": "USAID",
                "input": {"project": "Legacy input", "country": "Kenya"},
                "llm_mode": "true",
                "max_iterations": "0",
            },
        },
    )
    restored = store.get("job-inmem-1")
    assert restored is not None
    assert restored["state"]["donor_id"] == "usaid"
    assert restored["state"]["donor"] == "usaid"
    assert restored["state"]["input_context"]["project"] == "Legacy input"
    assert restored["state"]["input"]["country"] == "Kenya"
    assert restored["state"]["llm_mode"] is True
    assert restored["state"]["max_iterations"] == 1

    updated = store.update(
        "job-inmem-1",
        state={"donor": "EU", "input": {"project": "Digital governance", "country": "Moldova"}},
    )
    assert updated["state"]["donor_id"] == "eu"
    assert updated["state"]["donor"] == "eu"
    assert updated["state"]["input_context"]["project"] == "Digital governance"
    assert updated["state"]["input"]["country"] == "Moldova"


def test_sqlite_job_store_persists_json_and_rehydrates_strategy(tmp_path):
    db_path = tmp_path / "grantflow_state.db"
    store = SQLiteJobStore(str(db_path))
    strategy = DonorFactory.get_strategy("usaid")

    store.set(
        "job-1",
        {
            "status": "accepted",
            "hitl_enabled": False,
            "state": {
                "donor_id": "usaid",
                "donor": "usaid",
                "strategy": strategy,
                "donor_strategy": strategy,
                "input_context": {"project": "Health"},
            },
        },
    )

    restored = store.get("job-1")
    assert restored is not None
    assert restored["status"] == "accepted"
    assert restored["state"]["donor_id"] == "usaid"
    assert restored["state"]["strategy"].get_rag_collection() == "usaid_ads201"
    assert restored["state"]["donor_strategy"].get_rag_collection() == "usaid_ads201"

    with sqlite3.connect(str(db_path)) as conn:
        row = conn.execute("SELECT payload_json FROM jobs WHERE job_id = ?", ("job-1",)).fetchone()
    raw_payload = json.loads(row[0])
    assert "strategy" not in raw_payload["state"]
    assert "donor_strategy" not in raw_payload["state"]


def _make_redis_mock():
    store: dict = {}
    sets: dict = {}

    mock = MagicMock()
    mock.ping.return_value = True

    def setex(key, ttl, value):
        store[key] = value
        sets[key] = ttl

    def get(key):
        return store.get(key)

    def sadd(key, *values):
        store.setdefault(key, set())
        if isinstance(store[key], set):
            store[key].update(values)

    def smembers(key):
        val = store.get(key, set())
        return val if isinstance(val, set) else set()

    def expire(key, ttl):
        pass

    mock.setex.side_effect = setex
    mock.get.side_effect = get
    mock.sadd.side_effect = sadd
    mock.smembers.side_effect = smembers
    mock.expire.side_effect = expire
    return mock


def test_redis_job_store_set_get_update_list():
    mock_redis = _make_redis_mock()

    store = RedisJobStore()
    store._redis = mock_redis
    store._url = "redis://mock"

    payload = {"status": "accepted", "state": {"donor_id": "usaid", "input_context": {"project": "Test"}}}
    store.set("job-r1", payload)

    result = store.get("job-r1")
    assert result is not None
    assert result["status"] == "accepted"
    assert result["state"]["donor_id"] == "usaid"

    store.update("job-r1", status="running")
    updated = store.get("job-r1")
    assert updated["status"] == "running"

    all_jobs = store.list()
    assert "job-r1" in all_jobs
    assert all_jobs["job-r1"]["status"] == "running"


def test_redis_job_store_falls_back_to_inmem_when_redis_unavailable():
    store = RedisJobStore()
    store._url = None  # no Redis configured

    store.set("job-fallback", {"status": "accepted"})
    result = store.get("job-fallback")
    assert result is not None
    assert result["status"] == "accepted"

    all_jobs = store.list()
    assert "job-fallback" in all_jobs


def test_redis_job_store_falls_back_on_redis_error():
    mock_redis = MagicMock()
    mock_redis.setex.side_effect = Exception("connection lost")
    mock_redis.get.side_effect = Exception("connection lost")
    mock_redis.smembers.side_effect = Exception("connection lost")

    store = RedisJobStore()
    store._redis = mock_redis
    store._url = "redis://mock"

    store.set("job-err", {"status": "accepted"})
    result = store.get("job-err")
    assert result is not None
    assert result["status"] == "accepted"


def test_sqlite_job_store_get_restores_legacy_state_aliases(tmp_path):
    db_path = tmp_path / "grantflow_state.db"
    store = SQLiteJobStore(str(db_path))
    with sqlite3.connect(str(db_path)) as conn:
        conn.execute(
            """
            INSERT INTO jobs (job_id, payload_json, updated_at)
            VALUES (?, ?, CURRENT_TIMESTAMP)
            """,
            (
                "job-legacy",
                json.dumps(
                    {
                        "status": "accepted",
                        "state": {
                            "donor": "USAID",
                            "input": {"project": "Legacy state", "country": "Kenya"},
                            "llm_mode": "false",
                        },
                    }
                ),
            ),
        )

    restored = store.get("job-legacy")
    assert restored is not None
    assert restored["state"]["donor_id"] == "usaid"
    assert restored["state"]["donor"] == "usaid"
    assert restored["state"]["input_context"]["project"] == "Legacy state"
    assert restored["state"]["input"]["country"] == "Kenya"
    assert restored["state"]["strategy"].get_rag_collection() == "usaid_ads201"
    assert restored["state"]["donor_strategy"].get_rag_collection() == "usaid_ads201"


def test_sqlite_job_store_update_merges_and_restores(tmp_path):
    db_path = tmp_path / "grantflow_state.db"
    store = SQLiteJobStore(str(db_path))
    store.set("job-2", {"status": "accepted", "state": {"donor_id": "eu"}})

    updated = store.update("job-2", status="running", state={"donor_id": "eu", "progress": 50})
    assert updated["status"] == "running"
    assert updated["state"]["progress"] == 50
    assert updated["state"]["strategy"].get_rag_collection() == "eu_intpa"


def test_sqlite_job_store_list_returns_all_jobs(tmp_path):
    db_path = tmp_path / "grantflow_state.db"
    store = SQLiteJobStore(str(db_path))
    store.set("job-a", {"status": "accepted", "state": {"donor_id": "usaid"}})
    store.set("job-b", {"status": "done", "state": {"donor_id": "eu"}})

    items = store.list()
    assert set(items.keys()) == {"job-a", "job-b"}
    assert items["job-a"]["state"]["strategy"].get_rag_collection() == "usaid_ads201"
    assert items["job-b"]["state"]["strategy"].get_rag_collection() == "eu_intpa"


def test_sqlite_hitl_checkpoint_store_roundtrip(monkeypatch, tmp_path):
    monkeypatch.setenv("GRANTFLOW_HITL_STORE", "sqlite")
    monkeypatch.setenv("GRANTFLOW_SQLITE_PATH", str(tmp_path / "grantflow_state.db"))

    manager = HITLCheckpoint()
    strategy = DonorFactory.get_strategy("usaid")
    checkpoint_id = manager.create_checkpoint(
        stage="toc",
        state={"donor_id": "usaid", "strategy": strategy, "foo": "bar"},
        donor_id="usaid",
    )

    checkpoint = manager.get_checkpoint(checkpoint_id)
    assert checkpoint is not None
    assert checkpoint["status"] == HITLStatus.PENDING
    assert checkpoint["stage"] == "toc"
    assert checkpoint["state_snapshot"]["foo"] == "bar"
    assert "strategy" not in checkpoint["state_snapshot"]

    pending = manager.list_pending()
    assert any(cp["id"] == checkpoint_id for cp in pending)

    assert manager.approve(checkpoint_id, "ok")
    assert manager.is_approved(checkpoint_id) is True

    checkpoint = manager.get_checkpoint(checkpoint_id)
    assert checkpoint is not None
    assert checkpoint["status"] == HITLStatus.APPROVED
    assert checkpoint["feedback"] == "ok"


def test_sqlite_ingest_audit_store_roundtrip_and_filtering(tmp_path):
    db_path = tmp_path / "grantflow_state.db"
    store = SQLiteIngestAuditStore(str(db_path))

    store.append(
        {
            "event_id": "evt-1",
            "ts": "2026-02-25T10:00:00+00:00",
            "donor_id": "usaid",
            "namespace": "tenant_a/usaid_ads201",
            "filename": "ads.pdf",
            "content_type": "application/pdf",
            "metadata": {"doc_family": "donor_policy", "tenant_id": "tenant_a"},
            "result": {"chunks_ingested": 3},
        }
    )
    store.append(
        {
            "event_id": "evt-2",
            "ts": "2026-02-25T10:05:00+00:00",
            "donor_id": "eu",
            "namespace": "tenant_b/eu_intpa",
            "filename": "eu.pdf",
            "content_type": "application/pdf",
            "metadata": {"doc_family": "donor_policy", "tenant_id": "tenant_b"},
            "result": {"chunks_ingested": 2},
        }
    )
    store.append(
        {
            "event_id": "evt-3",
            "ts": "2026-02-25T10:10:00+00:00",
            "donor_id": "usaid",
            "namespace": "tenant_a/usaid_ads201",
            "filename": "kz-context.pdf",
            "content_type": "application/pdf",
            "metadata": {"doc_family": "country_context", "tenant_id": "tenant_a"},
            "result": {"chunks_ingested": 4},
        }
    )

    usaid_rows = store.list_recent(donor_id="usaid", limit=10)
    assert [row["filename"] for row in usaid_rows] == ["kz-context.pdf", "ads.pdf"]
    assert usaid_rows[0]["metadata"]["doc_family"] == "country_context"
    assert usaid_rows[1]["result"]["chunks_ingested"] == 3

    all_rows = store.list_recent(limit=2)
    assert len(all_rows) == 2
    assert [row["event_id"] for row in all_rows] == ["evt-3", "evt-2"]

    inventory = store.inventory(donor_id="usaid")
    assert len(inventory) == 2
    by_family = {row["doc_family"]: row for row in inventory}
    assert by_family["country_context"]["count"] == 1
    assert by_family["country_context"]["latest_filename"] == "kz-context.pdf"
    assert by_family["donor_policy"]["count"] == 1
    assert by_family["donor_policy"]["tenant_id"] == "tenant_a"

    tenant_rows = store.list_recent(limit=10, tenant_id="tenant_a")
    assert [row["event_id"] for row in tenant_rows] == ["evt-3", "evt-1"]
    tenant_inventory = store.inventory(tenant_id="tenant_a")
    assert len(tenant_inventory) == 2


def test_sqlite_stores_initialize_pragmas_and_schema_meta(monkeypatch, tmp_path):
    db_path = tmp_path / "grantflow_state.db"
    monkeypatch.setenv("GRANTFLOW_SQLITE_PATH", str(db_path))
    monkeypatch.setenv("GRANTFLOW_HITL_STORE", "sqlite")
    monkeypatch.setenv("GRANTFLOW_INGEST_STORE", "sqlite")
    monkeypatch.setenv("GRANTFLOW_SQLITE_BUSY_TIMEOUT_MS", "7000")

    SQLiteJobStore(str(db_path))
    HITLCheckpoint()
    SQLiteIngestAuditStore(str(db_path))

    with open_sqlite_connection(str(db_path)) as conn:
        journal_mode = conn.execute("PRAGMA journal_mode").fetchone()[0]
        busy_timeout = conn.execute("PRAGMA busy_timeout").fetchone()[0]
        rows = conn.execute("SELECT component, version FROM schema_meta ORDER BY component").fetchall()
        rows = [(row["component"], row["version"]) for row in rows]

    assert str(journal_mode).lower() == "wal"
    assert int(busy_timeout) == 7000
    assert ("hitl_checkpoints", 1) in rows
    assert ("ingest_audit", 1) in rows
    assert ("jobs", 1) in rows


def test_sqlite_job_store_upgrades_schema_meta_version_on_reinit(tmp_path):
    db_path = tmp_path / "grantflow_state.db"

    with sqlite3.connect(str(db_path)) as conn:
        conn.execute("""
            CREATE TABLE IF NOT EXISTS schema_meta (
              component TEXT PRIMARY KEY,
              version INTEGER NOT NULL,
              updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
            )
            """)
        conn.execute(
            "INSERT INTO schema_meta (component, version) VALUES (?, ?)",
            ("jobs", 0),
        )

    SQLiteJobStore(str(db_path))

    with sqlite3.connect(str(db_path)) as conn:
        version = conn.execute(
            "SELECT version FROM schema_meta WHERE component = ?",
            ("jobs",),
        ).fetchone()[0]

    assert int(version) == 1
