"""Integration tests for BenchmarkRunner (Issue #5)."""
import json
import shutil
import sqlite3
from pathlib import Path

import pytest

from gsql_track.bench import BenchmarkConfig, BenchmarkRunner, BenchmarkResults, ModelSpec
from gsql_track.config import TaskConfig
from tests.fixtures import DUMMY_FUNCTION_PATH, DUMMY_FUNCTION_FAILING_PATH

OUTPUT_DIR = Path("/tmp/gsql_test_bench_runner")
DB_PATH = OUTPUT_DIR / "benchmark.db"


@pytest.fixture(autouse=True, scope="module")
def setup_output():
    if OUTPUT_DIR.exists():
        shutil.rmtree(OUTPUT_DIR)
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    yield
    print(f"\n\nTest DB: gsql track serve --db {OUTPUT_DIR / 'track.db'}")


def _make_bench_config(n_seeds=2, n_workers=1, resume_mode="fresh", function=DUMMY_FUNCTION_PATH):
    return BenchmarkConfig(
        models=[
            ModelSpec(name="model_a", function=function, n_seeds=n_seeds,
                      train_config={"eval": {"primary_metric": "val.acc", "direction": "maximize"}}),
            ModelSpec(name="model_b", function=function, n_seeds=n_seeds,
                      train_config={"eval": {"primary_metric": "val.acc", "direction": "maximize"}}),
        ],
        tasks=[TaskConfig(name="task_x"), TaskConfig(name="task_y")],
        output=OUTPUT_DIR,
        n_workers=n_workers,
        show_progress=False,
        auto_confirm=True,
        resume_mode=resume_mode,
    )


def _query_db(sql, params=()):
    conn = sqlite3.connect(str(DB_PATH))
    conn.row_factory = sqlite3.Row
    rows = conn.execute(sql, params).fetchall()
    conn.close()
    return rows


# ─── CPU Tests ─────────────────────────────────────────────────────────────


class TestBenchSequential:
    def test_run_sequential(self):
        config = _make_bench_config(n_seeds=2)
        runner = BenchmarkRunner(config)
        results = runner.run_sequential()

        assert len(results.results) == 8
        assert set(results.df["model"].unique()) == {"model_a", "model_b"}
        assert set(results.df["task"].unique()) == {"task_x", "task_y"}

        # DONE files
        for model in ("model_a", "model_b"):
            for task in ("task_x", "task_y"):
                for seed in (1, 2):
                    assert (OUTPUT_DIR / model / task / f"seed_{seed}" / "DONE").exists()

        # track.db should exist alongside benchmark.db
        assert (OUTPUT_DIR / "track.db").exists()
        assert DB_PATH.exists()

    def test_db_benchmark_runs(self):
        """All 8 runs should be completed with correct model/task/seed combos."""
        rows = _query_db("SELECT model_name, task_name, seed, status FROM benchmark_runs")
        assert len(rows) == 8
        assert all(r["status"] == "completed" for r in rows)
        combos = {(r["model_name"], r["task_name"], r["seed"]) for r in rows}
        expected = {(m, t, s) for m in ("model_a", "model_b") for t in ("task_x", "task_y") for s in (1, 2)}
        assert combos == expected

    def test_db_benchmark_results(self):
        """Each run should have best_metrics with val.acc, test.acc, val.loss."""
        rows = _query_db("SELECT best_metrics FROM benchmark_results")
        assert len(rows) == 8
        for r in rows:
            metrics = json.loads(r["best_metrics"])
            assert "val.acc" in metrics
            assert "test.acc" in metrics
            assert "val.loss" in metrics

    def test_db_benchmark_state(self):
        """Benchmark state should have elapsed_time > 0."""
        rows = _query_db("SELECT key, value FROM benchmark_state WHERE key = 'elapsed_time'")
        assert len(rows) == 1
        assert float(json.loads(rows[0]["value"])) > 0


class TestBenchResumeMode:
    def test_resume_skips(self):
        runner = BenchmarkRunner(_make_bench_config(n_seeds=2))
        assert len(runner.run_sequential().results) == 8

        # Record end_times from first run
        rows_before = _query_db("SELECT id, end_time FROM benchmark_runs ORDER BY id")
        times_before = {r["id"]: r["end_time"] for r in rows_before}

        runner2 = BenchmarkRunner(_make_bench_config(n_seeds=2, resume_mode="resume"))
        assert len(runner2.run_sequential().results) == 8

        # End times should be unchanged — resume skipped all jobs
        rows_after = _query_db("SELECT id, end_time FROM benchmark_runs ORDER BY id")
        times_after = {r["id"]: r["end_time"] for r in rows_after}
        assert times_before == times_after


class TestBenchFreshMode:
    def test_fresh_reruns(self):
        runner = BenchmarkRunner(_make_bench_config(n_seeds=1))
        assert len(runner.run_sequential().results) == 4

        rows_before = _query_db("SELECT id, end_time FROM benchmark_runs ORDER BY id")
        times_before = {r["id"]: r["end_time"] for r in rows_before}

        runner2 = BenchmarkRunner(_make_bench_config(n_seeds=1, resume_mode="fresh"))
        assert len(runner2.run_sequential().results) == 4

        # Fresh mode should have re-executed — end_times should differ
        rows_after = _query_db("SELECT id, end_time FROM benchmark_runs ORDER BY id")
        times_after = {r["id"]: r["end_time"] for r in rows_after}
        assert any(times_after[rid] != times_before.get(rid) for rid in times_after)


# ─── Result Aggregation ───────────────────────────────────────────────────


class TestBenchResultAggregation:
    def test_aggregate(self):
        config = _make_bench_config(n_seeds=2)
        results = BenchmarkRunner(config).run_sequential()
        agg = results.aggregate(["val.acc"])

        assert len(agg) == 4  # 2 models × 2 tasks
        assert set(agg.columns) >= {"model", "task", "val.acc_mean", "val.acc_std"}
        assert set(agg["model"].unique()) == {"model_a", "model_b"}
        assert set(agg["task"].unique()) == {"task_x", "task_y"}
        # Means should be in valid range
        assert all(0 < v < 1 for v in agg["val.acc_mean"])

    def test_df_shape(self):
        config = _make_bench_config(n_seeds=2)
        results = BenchmarkRunner(config).run_sequential()
        assert results.df.shape[0] == 8
        assert "model" in results.df.columns
        assert "task" in results.df.columns
        assert "seed" in results.df.columns


# ─── Error Handling ───────────────────────────────────────────────────────


class TestBenchErrorHandling:
    def test_failed_job_recorded_in_db(self):
        """Runner raises on failure; DB should record the failed job with error."""
        config = _make_bench_config(n_seeds=3, function=DUMMY_FUNCTION_FAILING_PATH)
        runner = BenchmarkRunner(config)

        with pytest.raises(RuntimeError, match="Intentional failure"):
            runner.run_sequential()

        # The failed job should be recorded in DB with status='failed' and error populated
        failed = _query_db("SELECT model_name, task_name, seed, status, error FROM benchmark_runs WHERE status = 'failed'")
        assert len(failed) >= 1
        assert all("Intentional failure" in r["error"] for r in failed)

        # Jobs that ran before the failure should be completed
        completed = _query_db("SELECT status FROM benchmark_runs WHERE status = 'completed'")
        assert len(completed) >= 1  # at least some jobs ran before failure


# ─── GsqlTrack Integration ────────────────────────────────────────────────


class TestBenchGsqlTrackIntegration:
    def test_track_db_populated(self):
        """Verify track.db is populated via gsql log mode (bench default)."""
        # Run a fresh benchmark — the default log.mode='gsql' writes to track.db
        config = _make_bench_config(n_seeds=1)
        runner = BenchmarkRunner(config)
        results = runner.run_sequential()

        assert len(results.results) == 4
        track_db = OUTPUT_DIR / "track.db"
        assert track_db.exists()

        conn = sqlite3.connect(str(track_db))
        conn.row_factory = sqlite3.Row

        # Should have runs logged
        runs = conn.execute("SELECT * FROM runs").fetchall()
        assert len(runs) >= 4

        # Should have params logged
        params = conn.execute("SELECT * FROM params").fetchall()
        assert len(params) > 0

        conn.close()


# ─── GPU Variant ───────────────────────────────────────────────────────────


@pytest.mark.gpu
class TestBenchGPUParallel:
    def test_parallel_execution(self):
        # Clean slate — earlier tests used n_seeds=2, this uses n_seeds=3
        if OUTPUT_DIR.exists():
            shutil.rmtree(OUTPUT_DIR)
        OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
        config = BenchmarkConfig(
            models=[
                ModelSpec(name="model_a", function=DUMMY_FUNCTION_PATH, n_seeds=3,
                          train_config={"eval": {"primary_metric": "val.acc", "direction": "maximize"}}),
                ModelSpec(name="model_b", function=DUMMY_FUNCTION_PATH, n_seeds=3,
                          train_config={"eval": {"primary_metric": "val.acc", "direction": "maximize"}}),
            ],
            tasks=[TaskConfig(name="task_x"), TaskConfig(name="task_y")],
            output=OUTPUT_DIR,
            n_workers=4,
            device_ids=[0, 1],
            show_progress=False,
            auto_confirm=True,
            resume_mode="fresh",
        )
        results = BenchmarkRunner(config).run_parallel()
        assert len(results.results) == 12

    def test_parallel_db_state(self):
        """All 12 jobs should be completed with worker_id populated."""
        rows = _query_db("SELECT status, worker_id FROM benchmark_runs")
        completed = [r for r in rows if r["status"] == "completed"]
        assert len(completed) == 12
        # worker_id should be set for parallel runs
        assert all(r["worker_id"] is not None for r in completed)
