"""Integration tests for Tuner hyperparameter tuning (Issue #6)."""
import json
import shutil
import sqlite3
from pathlib import Path

import pytest

from gsql_track.tune import TuneConfig, Tuner, TuneResults, ModelTuneSpec, ScreenConfig, screen_sensitivity
from gsql_track.config import TaskConfig
from tests.fixtures import DUMMY_TUNE_FUNCTION_PATH, DUMMY_TUNE_FUNCTION_FAILING_PATH, DUMMY_TUNE_FUNCTION_MULTIHP_PATH

OUTPUT_DIR = Path("/tmp/gsql_test_tune")
DB_PATH = OUTPUT_DIR / "tuning.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_tune_config(n_trials=10, n_workers=1, resume_mode="fresh", function=DUMMY_TUNE_FUNCTION_PATH, **overrides):
    config_kwargs = dict(
        models=[
            ModelTuneSpec(
                name="model_a", function=function, n_seeds=1,
                train_config={
                    "lr": ["log", 0.001, 0.1],
                    "batch_size": ["choice", 8, 16, 32],
                    "eval": {"primary_metric": "val.acc", "direction": "maximize"},
                },
            ),
        ],
        tasks=[TaskConfig(name="task_x")],
        output=OUTPUT_DIR,
        n_trials=n_trials,
        metric="val.acc",
        direction="maximize",
        n_workers=n_workers,
        show_progress=False,
        auto_confirm=True,
        resume_mode=resume_mode,
    )
    config_kwargs.update(overrides)
    return TuneConfig(**config_kwargs)


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


# ─── Sequential Tests ─────────────────────────────────────────────────────


class TestTuneSequential:
    def test_run_sequential(self):
        config = _make_tune_config(n_trials=10)
        tuner = Tuner(config)
        results = tuner.run_sequential()

        assert len(results.results) == 1  # 1 model × 1 task
        r = results.results[0]
        assert r.model == "model_a"
        assert r.task == "task_x"
        assert r.n_trials == 10

    def test_db_tuning_jobs(self):
        """Tuning job should be completed with correct metadata."""
        rows = _query_db("SELECT model_name, task_name, status, best_value FROM tuning_jobs")
        assert len(rows) >= 1
        completed = [r for r in rows if r["status"] == "completed"]
        assert len(completed) >= 1
        assert all(r["best_value"] is not None for r in completed)

    def test_optuna_study_exists(self):
        """Optuna study DB should exist on disk."""
        study_path = OUTPUT_DIR / "studies" / "model_a_task_x.db"
        assert study_path.exists()

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


# ─── Resume / Fresh ───────────────────────────────────────────────────────


class TestTuneResumeMode:
    def test_resume_skips(self):
        """Resume mode should skip already-completed jobs."""
        if OUTPUT_DIR.exists():
            shutil.rmtree(OUTPUT_DIR)
        OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

        config = _make_tune_config(n_trials=5)
        Tuner(config).run_sequential()

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

        # Second run (resume) — should skip
        config2 = _make_tune_config(n_trials=5, resume_mode="resume")
        Tuner(config2).run_sequential()

        rows_after = _query_db("SELECT model_name, task_name, end_time FROM tuning_jobs ORDER BY model_name, task_name")
        times_after = {(r["model_name"], r["task_name"]): r["end_time"] for r in rows_after}
        assert times_before == times_after


class TestTuneFreshMode:
    def test_fresh_reruns(self):
        """Fresh mode should re-execute tuning jobs."""
        if OUTPUT_DIR.exists():
            shutil.rmtree(OUTPUT_DIR)
        OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

        config = _make_tune_config(n_trials=5)
        Tuner(config).run_sequential()

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

        config2 = _make_tune_config(n_trials=5, resume_mode="fresh")
        Tuner(config2).run_sequential()

        rows_after = _query_db("SELECT model_name, task_name, end_time FROM tuning_jobs ORDER BY model_name, task_name")
        times_after = {(r["model_name"], r["task_name"]): r["end_time"] for r in rows_after}
        assert any(times_after[k] != times_before.get(k) for k in times_after)


# ─── Best Config Export ──────────────────────────────────────────────────


class TestTuneBestConfigExport:
    def test_best_results_saved(self):
        """best_results directory should contain best trial info."""
        if OUTPUT_DIR.exists():
            shutil.rmtree(OUTPUT_DIR)
        OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

        config = _make_tune_config(n_trials=5)
        tuner = Tuner(config)
        results = tuner.run_sequential()
        tuner.save_results(results)

        best_results_dir = OUTPUT_DIR / "best_results"
        assert best_results_dir.exists()
        json_files = list(best_results_dir.glob("*.json"))
        assert len(json_files) >= 1
        data = json.loads(json_files[0].read_text())
        assert "params" in data
        assert "value" in data


# ─── Results API ─────────────────────────────────────────────────────────


class TestTuneResultsAPI:
    def test_get_best_config(self):
        """TuneResults.get_best_config should return best params."""
        if OUTPUT_DIR.exists():
            shutil.rmtree(OUTPUT_DIR)
        OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

        config = _make_tune_config(n_trials=5)
        results = Tuner(config).run_sequential()

        best = results.get_best_config("model_a", "task_x")
        assert best is not None

    def test_build_raw_dataframe(self):
        """_build_raw_dataframe should return trials with expected columns."""
        if OUTPUT_DIR.exists():
            shutil.rmtree(OUTPUT_DIR)
        OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

        config = _make_tune_config(n_trials=5)
        results = Tuner(config).run_sequential()

        df = results._build_raw_dataframe()
        assert len(df) == 5
        assert set(df.columns) >= {"model", "task", "trial", "primary_value"}


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


class TestTuneErrorHandling:
    def test_failed_fn_recorded_in_db(self):
        """Failed training function should result in DB error record."""
        if OUTPUT_DIR.exists():
            shutil.rmtree(OUTPUT_DIR)
        OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

        config = _make_tune_config(n_trials=3, function=DUMMY_TUNE_FUNCTION_FAILING_PATH)
        tuner = Tuner(config)

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

        failed = _query_db("SELECT model_name, task_name, status, error FROM tuning_jobs WHERE status = 'failed'")
        assert len(failed) >= 1
        assert all("Intentional tune failure" in (r["error"] or "") for r in failed)


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


class TestTuneGsqlTrackIntegration:
    def test_track_db_populated(self):
        """track.db should have runs with source='tune'."""
        if OUTPUT_DIR.exists():
            shutil.rmtree(OUTPUT_DIR)
        OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

        config = _make_tune_config(n_trials=5)
        Tuner(config).run_sequential()

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

        conn = sqlite3.connect(str(track_db))
        conn.row_factory = sqlite3.Row
        runs = conn.execute("SELECT * FROM runs WHERE source = 'tune'").fetchall()
        assert len(runs) >= 1
        params = conn.execute("SELECT * FROM params").fetchall()
        assert len(params) > 0
        conn.close()


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


@pytest.mark.gpu
class TestTuneGPUParallel:
    def test_parallel_execution(self):
        if OUTPUT_DIR.exists():
            shutil.rmtree(OUTPUT_DIR)
        OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

        config = TuneConfig(
            models=[
                ModelTuneSpec(
                    name="model_a", function=DUMMY_TUNE_FUNCTION_PATH, n_seeds=1,
                    train_config={
                        "lr": ["log", 0.001, 0.1],
                        "batch_size": ["choice", 8, 16, 32],
                        "eval": {"primary_metric": "val.acc", "direction": "maximize"},
                    },
                ),
                ModelTuneSpec(
                    name="model_b", function=DUMMY_TUNE_FUNCTION_PATH, n_seeds=1,
                    train_config={
                        "lr": ["log", 0.001, 0.1],
                        "batch_size": ["choice", 8, 16, 32],
                        "eval": {"primary_metric": "val.acc", "direction": "maximize"},
                    },
                ),
            ],
            tasks=[TaskConfig(name="task_x")],
            output=OUTPUT_DIR,
            n_trials=10,
            metric="val.acc",
            direction="maximize",
            n_workers=4,
            device_ids=[0, 1],
            show_progress=False,
            auto_confirm=True,
            resume_mode="fresh",
        )
        results = Tuner(config).run_parallel()
        assert len(results.results) == 2  # 2 models × 1 task

    def test_parallel_db_state(self):
        """Both jobs should be completed with worker_id populated."""
        rows = _query_db("SELECT status, worker_id FROM tuning_jobs")
        completed = [r for r in rows if r["status"] == "completed"]
        assert len(completed) == 2
        assert all(r["worker_id"] is not None for r in completed)


# ─── Screening Tests ─────────────────────────────────────────────────────

SCREEN_OUTPUT_DIR = Path("/tmp/gsql_test_screen")


def _make_multihp_tune_config(n_trials=50, **overrides):
    config_kwargs = dict(
        models=[
            ModelTuneSpec(
                name="model_hp", function=DUMMY_TUNE_FUNCTION_MULTIHP_PATH, n_seeds=1,
                train_config={
                    "lr": ["log", 0.0001, 0.1],
                    "weight_decay": ["log", 1e-6, 1e-2],
                    "hidden_dim": ["choice", 64, 128, 256, 512],
                    "dropout": ["range", 0.0, 0.5],
                    "n_layers": ["int", 1, 6],
                    "activation": ["choice", "relu", "gelu", "tanh"],
                    "warmup_steps": ["int", 10, 500],
                    "label_smoothing": ["range", 0.0, 0.2],
                    "eval": {"primary_metric": "val.acc", "direction": "maximize"},
                },
            ),
        ],
        tasks=[TaskConfig(name="task_s")],
        output=SCREEN_OUTPUT_DIR,
        n_trials=n_trials,
        metric="val.acc",
        direction="maximize",
        n_workers=1,
        show_progress=False,
        auto_confirm=True,
        resume_mode="fresh",
    )
    config_kwargs.update(overrides)
    return TuneConfig(**config_kwargs)


class TestScreenSensitivity:
    def test_screen_identifies_important_params(self):
        """Screening should rank lr and weight_decay as the most important params."""
        if SCREEN_OUTPUT_DIR.exists():
            shutil.rmtree(SCREEN_OUTPUT_DIR)
        SCREEN_OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

        config = _make_multihp_tune_config(n_trials=50)
        tuner = Tuner(config)
        job = tuner.jobs[0]
        sc = ScreenConfig(n_screen=40, evaluator="pedanova", threshold=0.05, show_plot=False)

        result = screen_sensitivity(job, config, sc)

        assert result.n_trials >= 40
        assert len(result.importances) > 0
        # lr and weight_decay should be in the top-ranked params
        sorted_params = sorted(result.importances.items(), key=lambda x: x[1], reverse=True)
        top_3_names = [p[0] for p in sorted_params[:3]]
        assert "lr" in top_3_names, f"lr not in top 3: {top_3_names}"

    def test_screen_reduces_space(self):
        """Screening should fix unimportant params and reduce the search space."""
        if SCREEN_OUTPUT_DIR.exists():
            shutil.rmtree(SCREEN_OUTPUT_DIR)
        SCREEN_OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

        config = _make_multihp_tune_config(n_trials=50)
        tuner = Tuner(config)
        job = tuner.jobs[0]
        sc = ScreenConfig(n_screen=40, evaluator="pedanova", threshold=0.05, show_plot=False)

        result = screen_sensitivity(job, config, sc)

        # Should have fewer params in reduced space than original
        assert len(result.important_params) < 8, f"Expected screening to identify < 8 important params, got {len(result.important_params)}"
        assert len(result.fixed_params) > 0, "Expected at least some params to be fixed"
        # Fixed params should have concrete values (not distribution specs)
        for param, val in result.fixed_params.items():
            assert not isinstance(val, list), f"Fixed param {param} should be a value, not {val}"

    def test_screen_warm_starts_tuning(self):
        """Screening trials should feed into full tuning (same study, warm-started)."""
        if SCREEN_OUTPUT_DIR.exists():
            shutil.rmtree(SCREEN_OUTPUT_DIR)
        SCREEN_OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

        config = _make_multihp_tune_config(
            n_trials=50,
            screen={"n_screen": 20, "evaluator": "pedanova", "threshold": 0.05, "show_plot": False},
        )
        tuner = Tuner(config)
        results = tuner.run_sequential()

        assert len(results.results) == 1
        r = results.results[0]
        # Total completed should be n_trials (screening + optimization)
        assert r.n_trials == 50, f"Expected 50 total trials, got {r.n_trials}"
