"""Tests for M5: Brain auto-tuner — closed-loop weight optimizer."""

from __future__ import annotations

import json

from radix_edge.meta.brain_autotuner import (
    normalize_weights,
    compute_job_metrics,
    identify_bottleneck,
    compute_weight_adjustments,
    validate_previous_adjustment,
    clamp,
    WEIGHT_KEYS,
    MAX_DELTA,
)


class TestNormalizeWeights:
    def test_basic(self):
        w = {"a": 0.2, "b": 0.3, "c": 0.5}
        n = normalize_weights(w)
        assert abs(sum(n.values()) - 1.0) < 1e-6

    def test_zero_total(self):
        w = {"a": 0.0, "b": 0.0}
        n = normalize_weights(w)
        assert n["a"] == 0.5
        assert n["b"] == 0.5

    def test_unnormalized(self):
        w = {"a": 2.0, "b": 3.0, "c": 5.0}
        n = normalize_weights(w)
        assert abs(n["a"] - 0.2) < 1e-6
        assert abs(n["b"] - 0.3) < 1e-6


class TestIdentifyBottleneck:
    def test_high_failure(self):
        assert identify_bottleneck({"failure_rate": 0.2}) == "node_pressure"

    def test_high_wait(self):
        assert identify_bottleneck({"failure_rate": 0.01, "median_wait_seconds": 500}) == "queue_depth"

    def test_low_gpu(self):
        assert identify_bottleneck({"failure_rate": 0.01, "median_wait_seconds": 10, "mean_gpu_utilization": 20.0}) == "gpu_saturation"

    def test_balanced(self):
        assert identify_bottleneck({"failure_rate": 0.01, "median_wait_seconds": 10, "mean_gpu_utilization": 60.0}) == "balanced"


class TestComputeWeightAdjustments:
    def test_boost_queue_depth(self):
        current = {"w_queue_depth": 0.5, "w_gpu_saturation": 0.3, "w_node_pressure": 0.2}
        metrics = {"median_wait_seconds": 600}

        new = compute_weight_adjustments(current, "queue_depth", metrics)

        # queue_depth weight should increase
        assert new["w_queue_depth"] > current["w_queue_depth"]
        # Sum should still be ~1.0
        total = sum(new[k] for k in WEIGHT_KEYS)
        assert abs(total - 1.0) < 1e-4

    def test_boost_gpu_saturation(self):
        current = {"w_queue_depth": 0.5, "w_gpu_saturation": 0.3, "w_node_pressure": 0.2}
        metrics = {"mean_gpu_utilization": 10.0}

        new = compute_weight_adjustments(current, "gpu_saturation", metrics)
        assert new["w_gpu_saturation"] > current["w_gpu_saturation"]

    def test_balanced_no_change(self):
        current = {"w_queue_depth": 0.5, "w_gpu_saturation": 0.3, "w_node_pressure": 0.2}
        new = compute_weight_adjustments(current, "balanced", {})
        assert new == current

    def test_max_delta_respected(self):
        current = {"w_queue_depth": 0.5, "w_gpu_saturation": 0.3, "w_node_pressure": 0.2}
        metrics = {"median_wait_seconds": 99999}  # Extreme severity

        new = compute_weight_adjustments(current, "queue_depth", metrics)
        # Change should not exceed MAX_DELTA (before normalization)
        # After normalization the values shift, but the intent is preserved
        assert new["w_queue_depth"] <= 0.9  # WEIGHT_BOUNDS max


class TestComputeJobMetrics:
    def test_no_jobs(self, db):
        metrics = compute_job_metrics(db)
        assert metrics["total_completed"] == 0
        assert metrics["failure_rate"] == 0.0

    def test_with_jobs(self, db):
        db.execute(
            """INSERT INTO jobs (job_id, user_id, image, status, runtime_seconds, created_at, started_at, completed_at)
            VALUES ('j1', 'u1', 'img', 'succeeded', 100.0, datetime('now', '-1 hour'), datetime('now', '-55 minutes'), datetime('now'))"""
        )
        db.execute(
            """INSERT INTO jobs (job_id, user_id, image, status, runtime_seconds, completed_at)
            VALUES ('j2', 'u1', 'img', 'failed', 10.0, datetime('now'))"""
        )
        db.commit()

        metrics = compute_job_metrics(db)
        assert metrics["total_completed"] == 1
        assert metrics["total_failed"] == 1
        assert metrics["failure_rate"] == 0.5


class TestValidatePreviousAdjustment:
    def test_no_audit(self, db):
        result = validate_previous_adjustment(db)
        assert result is None

    def test_revert_on_regression(self, db):
        # Insert old audit with good metrics
        db.execute(
            """INSERT INTO brain_audit (bottleneck, old_params, new_params, metrics)
            VALUES ('queue_depth', ?, ?, ?)""",
            (
                json.dumps({"w_queue_depth": 0.5, "w_gpu_saturation": 0.3, "w_node_pressure": 0.2}),
                json.dumps({"w_queue_depth": 0.6, "w_gpu_saturation": 0.25, "w_node_pressure": 0.15}),
                json.dumps({"failure_rate": 0.01, "median_wait_seconds": 10}),
            ),
        )
        db.commit()

        # Insert many jobs with high failure rate (regression)
        for i in range(15):
            db.execute(
                f"""INSERT INTO jobs (job_id, user_id, image, status, runtime_seconds, completed_at)
                VALUES ('j{i}', 'u1', 'img', 'failed', 10.0, datetime('now'))"""
            )
        db.commit()

        result = validate_previous_adjustment(db)
        assert result == "reverted"

        # Verify weights were reverted
        row = db.execute("SELECT * FROM brain_params WHERE key = 'current'").fetchone()
        assert row["w_queue_depth"] == 0.5
