"""Integration tests for Tracker and GsqlTrack lifecycle (Issue #4)."""
import concurrent.futures
import shutil
import sqlite3
from pathlib import Path

import pytest

from gsql_track.tracker import Tracker
from gsql_track.gsql import GsqlTrack
from tests.fixtures import make_train_config

OUTPUT_DIR = Path("/tmp/gsql_test_track")


@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'}")



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


class TestTrackerEpochBased:
    def test_tracker_epoch_based(self):
        config = make_train_config(OUTPUT_DIR / "epoch_based", tracking_mode="epochs", n_total=10)
        tracker = Tracker(config=config)

        for epoch in range(10):
            loss = 1.0 - (epoch * 0.08)
            acc = 0.5 + (epoch * 0.04)
            tracker.log_epoch({"val.acc": acc, "val.loss": loss}, epoch=epoch)
            if tracker.should_stop:
                break

        tracker.finalize()
        assert tracker.best_metrics and "val.acc" in tracker.best_metrics
        assert tracker.best_epoch >= 0


class TestTrackerStepBased:
    def test_tracker_step_based(self):
        config = make_train_config(
            OUTPUT_DIR / "step_based", tracking_mode="steps", n_total=100,
            primary_metric="val.loss", direction="minimize",
        )
        config.pbar.steps_per_epoch = 10
        tracker = Tracker(config=config)

        for step in range(0, 100, 10):
            loss = 1.0 - (step * 0.008)
            acc = 0.5 + (step * 0.004)
            tracker.log_step({"val.loss": loss, "val.acc": acc}, step=step)
            if tracker.should_stop:
                break

        tracker.finalize()
        assert tracker.best_metrics and "val.loss" in tracker.best_metrics


class TestTrackerEarlyStop:
    def test_early_stopping(self):
        config = make_train_config(OUTPUT_DIR / "early_stop", tracking_mode="epochs", n_total=20)
        tracker = Tracker(config=config)

        for epoch in range(20):
            acc = 0.5 + epoch * 0.05 if epoch < 5 else 0.5
            tracker.log_epoch({"val.acc": acc, "val.loss": 1.0 - acc}, epoch=epoch)
            if tracker.should_stop:
                break

        assert tracker.epoch < 20
        assert tracker.best_epoch <= 5


class TestTrackerEvaluateFunction:
    def test_evaluate_function(self):
        config = make_train_config(OUTPUT_DIR / "evaluate", tracking_mode="epochs", n_total=5)
        tracker = Tracker(config=config)
        call_count = [0]
        dummy_dataset = [1, 2, 3]

        def evaluate_fn(dataset=None, **kwargs):
            call_count[0] += 1
            return {"val.acc": 0.8 + call_count[0] * 0.01, "val.loss": 0.5}

        tracker.add_evaluate_function(evaluate_fn, dataset=dummy_dataset)
        for epoch in range(5):
            result = tracker.evaluate_data()
            tracker.log_epoch(result, epoch=epoch)

        tracker.finalize()
        assert call_count[0] == 5


class TestGsqlTrackAPI:
    def test_full_lifecycle(self):
        db_path = str(OUTPUT_DIR / "track.db")
        t = GsqlTrack("test-experiment", db_path=db_path)
        run = t.start_run("run-001")
        run.log_params({"lr": 0.001, "batch_size": 64, "model": "bert"})
        for step in range(5):
            run.log(step=step, loss=1.0 - step * 0.15, accuracy=0.5 + step * 0.1)
        run.log_predictions([
            {"label": "A", "prediction": "A", "confidence": 0.9, "text": "example 1"},
            {"label": "B", "prediction": "A", "confidence": 0.6, "text": "example 2"},
        ], text_key="text")
        run.finish()
        t.close()

        conn = sqlite3.connect(db_path)
        conn.row_factory = sqlite3.Row
        assert conn.execute("SELECT status FROM runs WHERE id = ?", (run.id,)).fetchone()["status"] == "FINISHED"
        params = {r["key"]: r["value"] for r in conn.execute("SELECT key, value FROM params WHERE run_id = ?", (run.id,)).fetchall()}
        assert params["lr"] == "0.001"
        assert conn.execute("SELECT COUNT(*) FROM metrics WHERE run_id = ?", (run.id,)).fetchone()[0] == 10
        assert conn.execute("SELECT COUNT(*) FROM predictions WHERE run_id = ?", (run.id,)).fetchone()[0] == 2
        conn.close()


# ─── Experiment Detail Page Data ───────────────────────────────────────────


def _create_seeded_run(db_path, experiment, run_name, params, target_acc, target_loss, n_epochs=20, seed=0, metrics=("val.acc", "val.loss", "val.f1")):
    """Helper: create a single run with smooth learning curves and seed-dependent noise."""
    import random
    rng = random.Random(seed * 100 + hash(run_name) % 1000)
    t = GsqlTrack(experiment, db_path=db_path)
    run = t.start_run(run_name)
    run.log_params(params)
    # Per-seed bias: shift entire curve so seeds are visually distinct (std ~0.04-0.06)
    seed_bias = (seed - 1) * 0.05 + rng.uniform(-0.02, 0.02)
    for epoch in range(n_epochs):
        progress = epoch / max(n_epochs - 1, 1)
        noise = rng.uniform(-0.015, 0.015)
        acc = 0.5 + (target_acc - 0.5) * (1 - (1 - progress) ** 2) + noise + seed_bias
        loss = 1.0 - (1.0 - target_loss) * (1 - (1 - progress) ** 2) + noise - seed_bias
        m = {}
        if "val.acc" in metrics: m["val.acc"] = acc
        if "val.loss" in metrics: m["val.loss"] = loss
        if "val.f1" in metrics: m["val.f1"] = acc * 0.97 + rng.uniform(-0.01, 0.01)
        run.log(step=epoch, **m)
    run.finish()
    t.close()
    return run.id


class TestExperimentWithSeedGroups:
    """Simulates a real benchmark: 2 models × 2 tasks × 3 seeds, plus a baseline."""

    def test_seed_groups(self):
        db_path = str(OUTPUT_DIR / "track.db")

        grid = {
            ("BERT_base", "SST2"): (0.82, 0.45),
            ("BERT_base", "MNLI"): (0.74, 0.58),
            ("BERT_large", "SST2"): (0.88, 0.35),
            ("BERT_large", "MNLI"): (0.80, 0.48),
        }

        all_run_ids = []
        for (model, task), (target_acc, target_loss) in grid.items():
            for seed in range(3):
                rid = _create_seeded_run(
                    db_path, "bert-benchmark", f"{model}/{task}/seed_{seed}",
                    params={"model": model, "task": task, "seed": str(seed), "lr": "2e-5", "batch_size": "32"},
                    target_acc=target_acc, target_loss=target_loss, seed=seed,
                )
                all_run_ids.append(rid)

        # Ungrouped baseline (no /seed_N suffix)
        _create_seeded_run(
            db_path, "bert-benchmark", "majority-baseline",
            params={"model": "majority", "task": "all"},
            target_acc=0.52, target_loss=0.69, metrics=("val.acc", "val.loss"),
        )

        conn = sqlite3.connect(db_path)
        exp_id = conn.execute("SELECT id FROM experiments WHERE name = 'bert-benchmark'").fetchone()[0]
        assert conn.execute("SELECT COUNT(*) FROM runs WHERE experiment_id = ?", (exp_id,)).fetchone()[0] == 13
        for rid in all_run_ids:
            assert conn.execute("SELECT COUNT(*) FROM metrics WHERE run_id = ?", (rid,)).fetchone()[0] == 60
        conn.close()


class TestMultiExperiment:
    """LR sweep experiment: 3 learning rates × 3 seeds, tests experiment filtering."""

    def test_separate_experiment(self):
        db_path = str(OUTPUT_DIR / "track.db")

        for lr_str, lr_val in [("1e-3", 0.001), ("1e-4", 0.0001), ("5e-5", 0.00005)]:
            for seed in range(3):
                _create_seeded_run(
                    db_path, "lr-sweep", f"lr_{lr_str}/seed_{seed}",
                    params={"lr": lr_str, "seed": str(seed)},
                    target_acc=0.75 + lr_val * 100, target_loss=0.5,
                    n_epochs=15, seed=seed, metrics=("val.acc", "val.loss"),
                )


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


@pytest.mark.gpu
class TestTrackGPUAsync:
    def test_concurrent_writes(self):
        db_path = str(OUTPUT_DIR / "track.db")

        def run_experiment(worker_id):
            t = GsqlTrack("concurrent-test", db_path=db_path)
            run = t.start_run(f"worker-{worker_id}")
            run.log_params({"worker": worker_id, "lr": 0.001 * worker_id})
            for step in range(10):
                run.log(step=step, loss=1.0 / (step + 1), acc=step * 0.1)
            run.finish()
            t.close()
            return run.id

        with concurrent.futures.ThreadPoolExecutor(max_workers=4) as pool:
            run_ids = [f.result() for f in [pool.submit(run_experiment, i) for i in range(4)]]

        assert len(set(run_ids)) == 4
        conn = sqlite3.connect(db_path)
        for rid in run_ids:
            assert conn.execute("SELECT COUNT(*) FROM metrics WHERE run_id = ?", (rid,)).fetchone()[0] == 20
        conn.close()
