"""Tests for sync pipeline — rebuild from live_entries, WAL checkpoint, invariants."""
from __future__ import annotations

import json
import os
import sqlite3
import tempfile
import unittest
from pathlib import Path

# Point data dir to a temp location BEFORE importing server
_tmpdir = tempfile.mkdtemp()
os.environ["MYBURN_DATA_DIR"] = _tmpdir

from myburn.server import (
    DB_PATH,
    _check_sync_integrity,
    _get_db,
    _init_db,
    _rebuild_daily_from_live,
    _rebuild_sessions_from_live,
    _wal_checkpoint,
    app,
)

from fastapi.testclient import TestClient


def _seed_live_entries(rows: list[tuple]):
    """Insert rows into live_entries. Each tuple:
    (timestamp, model, input_tokens, output_tokens,
     cache_read_tokens, cache_creation_tokens, cost, project, session_id)"""
    conn = _get_db()
    conn.executemany(
        "INSERT OR IGNORE INTO live_entries "
        "(timestamp, model, input_tokens, output_tokens, "
        "cache_read_tokens, cache_creation_tokens, cost, project, session_id) "
        "VALUES (?,?,?,?,?,?,?,?,?)",
        rows,
    )
    conn.commit()
    conn.close()


def _count_daily(conn, date=None, project=None):
    """Sum tokens from daily_usage, optionally filtered."""
    sql = "SELECT SUM(input_tokens) AS inp, SUM(output_tokens) AS out, SUM(cost) AS cost FROM daily_usage WHERE 1=1"
    params = []
    if date:
        sql += " AND date = ?"
        params.append(date)
    if project:
        sql += " AND project = ?"
        params.append(project)
    return conn.execute(sql, params).fetchone()


class TestRebuildDailyFromLive(unittest.TestCase):
    """Verify daily_usage is always consistent with live_entries."""

    def setUp(self):
        # Wipe DB for each test
        if DB_PATH.exists():
            DB_PATH.unlink()
        wal = Path(str(DB_PATH) + "-wal")
        shm = Path(str(DB_PATH) + "-shm")
        if wal.exists():
            wal.unlink()
        if shm.exists():
            shm.unlink()
        _init_db()

    def test_full_then_delta_preserves_totals(self):
        """Simulate full scan, then delta scan — totals must reflect ALL entries."""
        # Full scan: 3 entries on Feb 25
        full_entries = [
            ("2026-02-25T10:00:00Z", "claude-opus-4-6", 100000, 50000, 20000, 5000, 10.0, "proj-a", "s1"),
            ("2026-02-25T11:00:00Z", "claude-opus-4-6", 200000, 80000, 30000, 8000, 20.0, "proj-a", "s2"),
            ("2026-02-25T12:00:00Z", "claude-haiku-4-5", 50000, 10000, 5000, 1000, 1.0, "proj-b", "s3"),
        ]
        _seed_live_entries(full_entries)
        _rebuild_daily_from_live()

        conn = _get_db()
        r = _count_daily(conn, date="2026-02-25", project="__all__")
        self.assertEqual(r["inp"], 350000)
        self.assertEqual(r["out"], 140000)
        self.assertAlmostEqual(r["cost"], 31.0, places=2)

        # Delta scan: 1 more entry on Feb 25
        delta_entries = [
            ("2026-02-25T14:00:00Z", "claude-opus-4-6", 100000, 40000, 10000, 3000, 8.0, "proj-a", "s4"),
        ]
        _seed_live_entries(delta_entries)
        _rebuild_daily_from_live()

        # Totals must be full + delta (not just delta)
        r2 = _count_daily(conn, date="2026-02-25", project="__all__")
        self.assertEqual(r2["inp"], 450000)
        self.assertEqual(r2["out"], 180000)
        self.assertAlmostEqual(r2["cost"], 39.0, places=2)
        conn.close()

    def test_partial_scan_cannot_overwrite(self):
        """Even if only 10 of 110 entries come from delta, daily_usage has all 110."""
        # Insert 100 entries
        entries = [
            (f"2026-02-20T{10 + i // 60:02d}:{i % 60:02d}:00Z", "claude-opus-4-6",
             1000, 500, 200, 100, 0.5, "proj-a", f"s{i}")
            for i in range(100)
        ]
        _seed_live_entries(entries)
        _rebuild_daily_from_live()

        conn = _get_db()
        r1 = _count_daily(conn, date="2026-02-20", project="__all__")
        self.assertEqual(r1["inp"], 100000)

        # Insert 10 more entries (simulating delta)
        more = [
            (f"2026-02-20T{12 + i // 60:02d}:{40 + i % 60:02d}:00Z", "claude-opus-4-6",
             1000, 500, 200, 100, 0.5, "proj-a", f"s{100 + i}")
            for i in range(10)
        ]
        _seed_live_entries(more)
        _rebuild_daily_from_live()

        r2 = _count_daily(conn, date="2026-02-20", project="__all__")
        self.assertEqual(r2["inp"], 110000)  # All 110, not just 10
        conn.close()

    def test_cleanup_preserves_old_daily(self):
        """After cleanup deletes old live_entries, daily_usage for those dates persists."""
        # Insert entries for recent day and old day
        _seed_live_entries([
            ("2026-02-25T10:00:00Z", "claude-opus-4-6", 100000, 50000, 20000, 5000, 10.0, "proj-a", "s1"),
            ("2026-01-01T10:00:00Z", "claude-opus-4-6", 80000, 40000, 15000, 4000, 8.0, "proj-a", "s2"),
        ])
        _rebuild_daily_from_live()

        conn = _get_db()
        # Both dates have daily_usage
        old = _count_daily(conn, date="2026-01-01", project="__all__")
        self.assertEqual(old["inp"], 80000)

        # Simulate cleanup: delete old live_entries (like _cleanup_old_live_entries)
        conn.execute("DELETE FROM live_entries WHERE timestamp < '2026-02-01'")
        conn.commit()
        conn.close()

        # Rebuild — should only touch dates still in live_entries
        _rebuild_daily_from_live()

        conn = _get_db()
        # Old daily_usage preserved (not deleted)
        old2 = _count_daily(conn, date="2026-01-01", project="__all__")
        self.assertEqual(old2["inp"], 80000)

        # Recent day still correct
        recent = _count_daily(conn, date="2026-02-25", project="__all__")
        self.assertEqual(recent["inp"], 100000)
        conn.close()

    def test_cost_recalc_propagates(self):
        """After live_entries costs are updated, rebuild propagates them."""
        _seed_live_entries([
            ("2026-02-25T10:00:00Z", "claude-opus-4-6", 100000, 50000, 20000, 5000, 10.0, "proj-a", "s1"),
        ])
        _rebuild_daily_from_live()

        conn = _get_db()
        r1 = _count_daily(conn, date="2026-02-25", project="__all__")
        self.assertAlmostEqual(r1["cost"], 10.0, places=2)

        # Simulate cost recalculation
        conn.execute("UPDATE live_entries SET cost = 15.0 WHERE timestamp = '2026-02-25T10:00:00Z'")
        conn.commit()
        conn.close()

        _rebuild_daily_from_live()

        conn = _get_db()
        r2 = _count_daily(conn, date="2026-02-25", project="__all__")
        self.assertAlmostEqual(r2["cost"], 15.0, places=2)
        conn.close()

    def test_all_aggregate_rows_match_per_project(self):
        """__all__ rows must equal sum of per-project rows."""
        _seed_live_entries([
            ("2026-02-25T10:00:00Z", "claude-opus-4-6", 100000, 50000, 20000, 5000, 10.0, "proj-a", "s1"),
            ("2026-02-25T11:00:00Z", "claude-opus-4-6", 60000, 30000, 10000, 3000, 6.0, "proj-b", "s2"),
            ("2026-02-25T12:00:00Z", "claude-haiku-4-5", 20000, 5000, 2000, 500, 0.5, "proj-a", "s3"),
        ])
        _rebuild_daily_from_live()

        conn = _get_db()
        # Sum all per-project rows (not __all__)
        per_project = conn.execute(
            "SELECT SUM(input_tokens) AS inp FROM daily_usage WHERE project != '__all__' AND date = '2026-02-25'"
        ).fetchone()
        # __all__ rows
        all_row = _count_daily(conn, date="2026-02-25", project="__all__")
        self.assertEqual(per_project["inp"], all_row["inp"])
        conn.close()


class TestRebuildSessionsFromLive(unittest.TestCase):
    """Verify sessions are consistent with live_entries."""

    def setUp(self):
        if DB_PATH.exists():
            DB_PATH.unlink()
        wal = Path(str(DB_PATH) + "-wal")
        shm = Path(str(DB_PATH) + "-shm")
        if wal.exists():
            wal.unlink()
        if shm.exists():
            shm.unlink()
        _init_db()

    def test_sessions_consistent_with_live(self):
        """Session totals must match live_entries grouped by project."""
        _seed_live_entries([
            ("2026-02-25T10:00:00Z", "claude-opus-4-6", 100000, 50000, 20000, 5000, 10.0, "proj-a", "s1"),
            ("2026-02-25T11:00:00Z", "claude-opus-4-6", 60000, 30000, 10000, 3000, 6.0, "proj-a", "s2"),
            ("2026-02-25T12:00:00Z", "claude-haiku-4-5", 20000, 5000, 2000, 500, 0.5, "proj-b", "s3"),
        ])
        _rebuild_sessions_from_live()

        conn = _get_db()
        sessions = conn.execute(
            "SELECT session_id, input_tokens, output_tokens, total_tokens, total_cost "
            "FROM sessions ORDER BY session_id"
        ).fetchall()
        self.assertEqual(len(sessions), 2)

        # proj-a: 160000 in, 80000 out
        sa = [s for s in sessions if s["session_id"] == "proj-a"][0]
        self.assertEqual(sa["input_tokens"], 160000)
        self.assertEqual(sa["output_tokens"], 80000)
        self.assertAlmostEqual(sa["total_cost"], 16.0, places=2)

        # proj-b: 20000 in, 5000 out
        sb = [s for s in sessions if s["session_id"] == "proj-b"][0]
        self.assertEqual(sb["input_tokens"], 20000)
        self.assertEqual(sb["output_tokens"], 5000)
        conn.close()

    def test_sessions_have_model_breakdowns(self):
        """Sessions must include models_used and model_breakdowns JSON."""
        _seed_live_entries([
            ("2026-02-25T10:00:00Z", "claude-opus-4-6", 100000, 50000, 20000, 5000, 10.0, "proj-a", "s1"),
            ("2026-02-25T11:00:00Z", "claude-haiku-4-5", 20000, 5000, 2000, 500, 0.5, "proj-a", "s2"),
        ])
        _rebuild_sessions_from_live()

        conn = _get_db()
        row = conn.execute(
            "SELECT models_used, model_breakdowns FROM sessions WHERE session_id = 'proj-a'"
        ).fetchone()
        models = json.loads(row["models_used"])
        breakdowns = json.loads(row["model_breakdowns"])
        self.assertIn("claude-opus-4-6", models)
        self.assertIn("claude-haiku-4-5", models)
        self.assertEqual(len(breakdowns), 2)

        # Check breakdown structure
        opus_bd = [b for b in breakdowns if b["modelName"] == "claude-opus-4-6"][0]
        self.assertEqual(opus_bd["inputTokens"], 100000)
        conn.close()


class TestWALCheckpoint(unittest.TestCase):

    def setUp(self):
        if DB_PATH.exists():
            DB_PATH.unlink()
        wal = Path(str(DB_PATH) + "-wal")
        shm = Path(str(DB_PATH) + "-shm")
        if wal.exists():
            wal.unlink()
        if shm.exists():
            shm.unlink()
        _init_db()

    def test_db_file_nonzero_after_checkpoint(self):
        """After writing data and checkpointing, DB file must be > 0 bytes."""
        _seed_live_entries([
            ("2026-02-25T10:00:00Z", "claude-opus-4-6", 100000, 50000, 20000, 5000, 10.0, "proj-a", "s1"),
        ])
        _wal_checkpoint()
        self.assertGreater(DB_PATH.stat().st_size, 0)


class TestSyncIntegrity(unittest.TestCase):

    def setUp(self):
        if DB_PATH.exists():
            DB_PATH.unlink()
        wal = Path(str(DB_PATH) + "-wal")
        shm = Path(str(DB_PATH) + "-shm")
        if wal.exists():
            wal.unlink()
        if shm.exists():
            shm.unlink()
        _init_db()

    def test_live_entries_append_only(self):
        """Two inserts with overlapping entries — INSERT OR IGNORE, no duplicates."""
        entries = [
            ("2026-02-25T10:00:00Z", "claude-opus-4-6", 100000, 50000, 20000, 5000, 10.0, "proj-a", "s1"),
            ("2026-02-25T11:00:00Z", "claude-opus-4-6", 60000, 30000, 10000, 3000, 6.0, "proj-a", "s2"),
        ]
        _seed_live_entries(entries)

        # Re-insert same entries + 1 new
        overlap = entries + [
            ("2026-02-25T12:00:00Z", "claude-haiku-4-5", 20000, 5000, 2000, 500, 0.5, "proj-b", "s3"),
        ]
        _seed_live_entries(overlap)

        conn = _get_db()
        count = conn.execute("SELECT COUNT(*) AS c FROM live_entries").fetchone()["c"]
        self.assertEqual(count, 3)  # 3 unique, not 5
        conn.close()

    def test_empty_live_entries_rebuild_is_noop(self):
        """Rebuilding with no live_entries should not crash or delete historical data."""
        # Seed daily_usage directly (historical)
        conn = _get_db()
        conn.execute(
            "INSERT INTO daily_usage (date, project, model_name, input_tokens, output_tokens, "
            "cache_creation_tokens, cache_read_tokens, cost) VALUES (?,?,?,?,?,?,?,?)",
            ("2025-12-01", "__all__", "claude-opus-4-6", 50000, 20000, 5000, 2000, 5.0),
        )
        conn.commit()
        conn.close()

        # Rebuild with empty live_entries
        _rebuild_daily_from_live()

        # Historical row must survive
        conn = _get_db()
        r = _count_daily(conn, date="2025-12-01", project="__all__")
        self.assertEqual(r["inp"], 50000)
        conn.close()

    def test_multiple_projects_same_date(self):
        """Multiple projects on same date produce correct per-project and __all__ rows."""
        _seed_live_entries([
            ("2026-02-25T10:00:00Z", "claude-opus-4-6", 100000, 50000, 20000, 5000, 10.0, "proj-a", "s1"),
            ("2026-02-25T11:00:00Z", "claude-opus-4-6", 60000, 30000, 10000, 3000, 6.0, "proj-b", "s2"),
        ])
        _rebuild_daily_from_live()

        conn = _get_db()
        pa = _count_daily(conn, date="2026-02-25", project="proj-a")
        pb = _count_daily(conn, date="2026-02-25", project="proj-b")
        all_ = _count_daily(conn, date="2026-02-25", project="__all__")

        self.assertEqual(pa["inp"], 100000)
        self.assertEqual(pb["inp"], 60000)
        self.assertEqual(all_["inp"], 160000)
        conn.close()

    def test_rebuild_is_idempotent(self):
        """Running rebuild twice produces identical results."""
        _seed_live_entries([
            ("2026-02-25T10:00:00Z", "claude-opus-4-6", 100000, 50000, 20000, 5000, 10.0, "proj-a", "s1"),
        ])
        _rebuild_daily_from_live()

        conn = _get_db()
        r1 = _count_daily(conn, date="2026-02-25", project="__all__")
        conn.close()

        _rebuild_daily_from_live()

        conn = _get_db()
        r2 = _count_daily(conn, date="2026-02-25", project="__all__")
        self.assertEqual(r1["inp"], r2["inp"])
        self.assertAlmostEqual(r1["cost"], r2["cost"], places=4)

        # No duplicate rows
        count = conn.execute(
            "SELECT COUNT(*) AS c FROM daily_usage WHERE date = '2026-02-25' AND project = '__all__'"
        ).fetchone()["c"]
        # Should be exactly 1 model (claude-opus-4-6)
        self.assertEqual(count, 1)
        conn.close()


class TestSyncIntegrityCheck(unittest.TestCase):
    """Verify _check_sync_integrity detects and corrects drift."""

    def setUp(self):
        if DB_PATH.exists():
            DB_PATH.unlink()
        wal = Path(str(DB_PATH) + "-wal")
        shm = Path(str(DB_PATH) + "-shm")
        if wal.exists():
            wal.unlink()
        if shm.exists():
            shm.unlink()
        _init_db()

    def test_integrity_check_consistent(self):
        """After normal rebuild, integrity check returns True."""
        _seed_live_entries([
            ("2026-02-25T10:00:00Z", "claude-opus-4-6", 100000, 50000, 20000, 5000, 10.0, "proj-a", "s1"),
            ("2026-02-25T11:00:00Z", "claude-haiku-4-5", 20000, 5000, 2000, 500, 0.5, "proj-b", "s2"),
        ])
        _rebuild_daily_from_live()
        self.assertTrue(_check_sync_integrity())

    def test_integrity_check_detects_drift(self):
        """If daily_usage diverges from live_entries, integrity check returns False."""
        _seed_live_entries([
            ("2026-02-25T10:00:00Z", "claude-opus-4-6", 100000, 50000, 20000, 5000, 10.0, "proj-a", "s1"),
        ])
        _rebuild_daily_from_live()

        # Manually corrupt daily_usage to create drift
        conn = _get_db()
        conn.execute("UPDATE daily_usage SET input_tokens = 999 WHERE project = '__all__'")
        conn.commit()
        conn.close()

        self.assertFalse(_check_sync_integrity())

    def test_integrity_drift_triggers_rebuild(self):
        """After detecting drift, a corrective rebuild restores consistency."""
        _seed_live_entries([
            ("2026-02-25T10:00:00Z", "claude-opus-4-6", 100000, 50000, 20000, 5000, 10.0, "proj-a", "s1"),
        ])
        _rebuild_daily_from_live()

        # Corrupt daily_usage
        conn = _get_db()
        conn.execute("UPDATE daily_usage SET input_tokens = 999 WHERE project = '__all__'")
        conn.commit()
        conn.close()

        self.assertFalse(_check_sync_integrity())

        # Corrective rebuild (same pattern as _native_sync)
        _rebuild_daily_from_live()
        self.assertTrue(_check_sync_integrity())

    def test_integrity_empty_db(self):
        """Empty DB (no live_entries) should return True (no drift possible)."""
        self.assertTrue(_check_sync_integrity())


class TestHealthEndpoint(unittest.TestCase):
    """Verify /api/health returns expected fields."""

    def setUp(self):
        if DB_PATH.exists():
            DB_PATH.unlink()
        wal = Path(str(DB_PATH) + "-wal")
        shm = Path(str(DB_PATH) + "-shm")
        if wal.exists():
            wal.unlink()
        if shm.exists():
            shm.unlink()
        _init_db()
        # Mark DB as ready for the test client
        from myburn.server import _db_ready
        _db_ready.set()

    def test_health_endpoint_fields(self):
        """GET /api/health returns all expected fields."""
        _seed_live_entries([
            ("2026-02-25T10:00:00Z", "claude-opus-4-6", 100000, 50000, 20000, 5000, 10.0, "proj-a", "s1"),
        ])
        _rebuild_daily_from_live()

        client = TestClient(app)
        resp = client.get("/api/health")
        self.assertEqual(resp.status_code, 200)

        data = resp.json()
        expected_keys = {"status", "db_ready", "live_entries", "daily_rows",
                         "live_tokens", "daily_tokens", "integrity", "db_bytes"}
        self.assertTrue(expected_keys.issubset(set(data.keys())))
        self.assertEqual(data["status"], "ok")
        self.assertTrue(data["integrity"])
        self.assertEqual(data["live_entries"], 1)
        self.assertGreater(data["live_tokens"], 0)
        self.assertGreater(data["db_bytes"], 0)

    def test_health_reports_drift(self):
        """If daily_usage is corrupted, health reports drift status."""
        _seed_live_entries([
            ("2026-02-25T10:00:00Z", "claude-opus-4-6", 100000, 50000, 20000, 5000, 10.0, "proj-a", "s1"),
        ])
        _rebuild_daily_from_live()

        # Corrupt
        conn = _get_db()
        conn.execute("UPDATE daily_usage SET input_tokens = 1 WHERE project = '__all__'")
        conn.commit()
        conn.close()

        client = TestClient(app)
        resp = client.get("/api/health")
        data = resp.json()
        self.assertEqual(data["status"], "drift")
        self.assertFalse(data["integrity"])


if __name__ == "__main__":
    unittest.main()
