"""Brain auto-tuner — closed-loop scheduling weight optimizer.

Ported from services/cloud-api/brain_autotuner.py — runs locally every 60s
instead of hourly via Lambda.
"""

from __future__ import annotations

import asyncio
import json
import logging
import statistics
from datetime import datetime, timezone
from typing import Optional

from radix_edge.config import get_config
from radix_edge.db import get_db

logger = logging.getLogger("radix_edge.meta")

# Guard rails
MIN_COMPLETED_JOBS = 10
MAX_DELTA = 0.05
WEIGHT_KEYS = ("w_queue_depth", "w_gpu_saturation", "w_node_pressure")
WEIGHT_BOUNDS = (0.05, 0.90)
FAILURE_RATE_REGRESSION = 0.05  # absolute increase triggers revert
WAIT_TIME_REGRESSION = 0.20  # 20% relative increase triggers revert


def clamp(value: float, lo: float, hi: float) -> float:
    return max(lo, min(hi, value))


def normalize_weights(weights: dict[str, float]) -> dict[str, float]:
    """Normalize weights to sum to 1.0."""
    total = sum(weights.values())
    if total <= 0:
        n = len(weights)
        return {k: 1.0 / n for k in weights}
    return {k: v / total for k, v in weights.items()}


def compute_job_metrics(db_conn, hours: int = 2) -> dict:
    """Compute aggregate metrics from recently completed jobs."""
    rows = db_conn.execute(
        """SELECT status, runtime_seconds, created_at, started_at
        FROM jobs
        WHERE status IN ('succeeded', 'failed')
        AND completed_at > datetime('now', ? || ' hours')""",
        (f"-{hours}",),
    ).fetchall()

    if not rows:
        return {
            "median_wait_seconds": 0.0,
            "p95_wait_seconds": 0.0,
            "mean_gpu_utilization": 0.0,
            "failure_rate": 0.0,
            "total_completed": 0,
            "total_failed": 0,
        }

    completed = [r for r in rows if r["status"] == "succeeded"]
    failed = [r for r in rows if r["status"] == "failed"]

    # Wait times
    wait_times = []
    for r in rows:
        if r["created_at"] and r["started_at"]:
            try:
                t0 = datetime.fromisoformat(r["created_at"].replace("Z", "+00:00"))
                t1 = datetime.fromisoformat(r["started_at"].replace("Z", "+00:00"))
                wait_s = (t1 - t0).total_seconds()
                if wait_s >= 0:
                    wait_times.append(wait_s)
            except (ValueError, TypeError):
                pass

    if wait_times:
        median_wait = statistics.median(wait_times)
        sorted_waits = sorted(wait_times)
        p95_idx = int(len(sorted_waits) * 0.95)
        p95_wait = sorted_waits[min(p95_idx, len(sorted_waits) - 1)]
    else:
        median_wait = 0.0
        p95_wait = 0.0

    # GPU utilization from gpu table (current snapshot)
    gpu_rows = db_conn.execute("SELECT utilization_pct FROM gpus").fetchall()
    mean_gpu = statistics.mean([r["utilization_pct"] for r in gpu_rows]) if gpu_rows else 0.0

    total = len(rows)
    failure_rate = len(failed) / total if total > 0 else 0.0

    return {
        "median_wait_seconds": median_wait,
        "p95_wait_seconds": p95_wait,
        "mean_gpu_utilization": mean_gpu,
        "failure_rate": failure_rate,
        "total_completed": len(completed),
        "total_failed": len(failed),
    }


def identify_bottleneck(metrics: dict) -> str:
    """Identify dominant scheduling bottleneck.

    Priority: failure → queue wait → GPU underuse → balanced.
    """
    if metrics.get("failure_rate", 0) > 0.15:
        return "node_pressure"
    if metrics.get("median_wait_seconds", 0) > 300:
        return "queue_depth"
    if 0 < metrics.get("mean_gpu_utilization", 1.0) < 40.0:
        return "gpu_saturation"
    return "balanced"


def compute_weight_adjustments(current_weights: dict[str, float], bottleneck: str,
                                metrics: dict) -> dict[str, float]:
    """Compute new weights by boosting the bottleneck dimension."""
    boost_map = {
        "queue_depth": "w_queue_depth",
        "gpu_saturation": "w_gpu_saturation",
        "node_pressure": "w_node_pressure",
    }
    boost_key = boost_map.get(bottleneck)
    if not boost_key:
        return dict(current_weights)

    # Severity-proportional delta
    if bottleneck == "node_pressure":
        severity = min(metrics.get("failure_rate", 0) / 0.3, 1.0)
    elif bottleneck == "queue_depth":
        severity = min(metrics.get("median_wait_seconds", 0) / 600, 1.0)
    else:
        gpu_util = metrics.get("mean_gpu_utilization", 40.0)
        severity = min((40.0 - gpu_util) / 40.0, 1.0) if gpu_util < 40.0 else 0.0

    delta = MAX_DELTA * severity

    new_weights = dict(current_weights)
    new_weights[boost_key] = new_weights.get(boost_key, 0.0) + delta

    # Clamp and normalize twice (clamping can break sum)
    for _ in range(2):
        for k in WEIGHT_KEYS:
            new_weights[k] = clamp(new_weights.get(k, 0.0), *WEIGHT_BOUNDS)
        subset = {k: new_weights[k] for k in WEIGHT_KEYS}
        normalized = normalize_weights(subset)
        for k in WEIGHT_KEYS:
            new_weights[k] = round(normalized[k], 6)

    return new_weights


def validate_previous_adjustment(db_conn) -> Optional[str]:
    """Check if the last adjustment regressed metrics. Auto-revert if so.

    Returns 'validated', 'reverted', or None.
    """
    audit = db_conn.execute(
        "SELECT * FROM brain_audit ORDER BY created_at DESC LIMIT 1"
    ).fetchone()
    if not audit:
        return None

    old_metrics = json.loads(audit["metrics"]) if audit["metrics"] else {}
    if not old_metrics:
        return None

    config = get_config()
    metrics = compute_job_metrics(db_conn, hours=2)
    total = metrics["total_completed"] + metrics["total_failed"]
    if total < config.autotuner_min_jobs:
        return None

    # Check for regression
    old_failure = float(old_metrics.get("failure_rate", 0))
    new_failure = float(metrics.get("failure_rate", 0))
    failure_regressed = (new_failure - old_failure) > FAILURE_RATE_REGRESSION

    old_wait = float(old_metrics.get("median_wait_seconds", 0))
    new_wait = float(metrics.get("median_wait_seconds", 0))
    if old_wait > 0:
        wait_regressed = (new_wait - old_wait) / old_wait > WAIT_TIME_REGRESSION
    else:
        wait_regressed = new_wait > 300

    if failure_regressed or wait_regressed:
        # Revert to old params
        old_params = json.loads(audit["old_params"]) if audit["old_params"] else {}
        if old_params:
            now = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%fZ")
            for k, v in old_params.items():
                if k in WEIGHT_KEYS:
                    db_conn.execute(
                        f"UPDATE brain_params SET {k} = ?, updated_at = ? WHERE key = 'current'",
                        (v, now),
                    )
            db_conn.commit()
            logger.warning("Auto-tuner REVERTED previous adjustment (regression detected)")
        return "reverted"

    return "validated"


# --- Background loop ---

async def autotuner_loop():
    """Background task: continuously tune scheduling weights."""
    config = get_config()
    if not config.autotuner_enabled:
        logger.info("Brain auto-tuner disabled")
        return

    logger.info("Brain auto-tuner started (interval=%.0fs)", config.autotuner_interval)

    while True:
        try:
            with get_db() as db:
                # Validate previous adjustment first
                validation = validate_previous_adjustment(db)
                if validation == "reverted":
                    await asyncio.sleep(config.autotuner_interval)
                    continue

                # Compute current metrics
                metrics = compute_job_metrics(db, hours=2)
                total = metrics["total_completed"] + metrics["total_failed"]
                if total < config.autotuner_min_jobs:
                    await asyncio.sleep(config.autotuner_interval)
                    continue

                # Identify bottleneck
                bottleneck = identify_bottleneck(metrics)
                if bottleneck == "balanced":
                    await asyncio.sleep(config.autotuner_interval)
                    continue

                # Get current weights
                row = db.execute("SELECT * FROM brain_params WHERE key = 'current'").fetchone()
                current = {k: row[k] for k in WEIGHT_KEYS}

                # Compute new weights
                new_weights = compute_weight_adjustments(current, bottleneck, metrics)

                # Check if change is meaningful
                max_change = max(abs(new_weights[k] - current[k]) for k in WEIGHT_KEYS)
                if max_change < 0.001:
                    await asyncio.sleep(config.autotuner_interval)
                    continue

                # Apply new weights
                now = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%fZ")
                for k in WEIGHT_KEYS:
                    db.execute(
                        f"UPDATE brain_params SET {k} = ?, updated_at = ? WHERE key = 'current'",
                        (new_weights[k], now),
                    )

                # Record audit
                db.execute(
                    """INSERT INTO brain_audit (bottleneck, old_params, new_params, metrics)
                    VALUES (?, ?, ?, ?)""",
                    (
                        bottleneck,
                        json.dumps({k: round(v, 6) for k, v in current.items()}),
                        json.dumps({k: round(v, 6) for k, v in new_weights.items()}),
                        json.dumps({k: round(v, 4) if isinstance(v, float) else v for k, v in metrics.items()}),
                    ),
                )

                # Prune old audit entries (30 days)
                db.execute("DELETE FROM brain_audit WHERE created_at < datetime('now', '-30 days')")

                logger.info(
                    "Auto-tuner: bottleneck=%s, adjusted %s",
                    bottleneck,
                    {k: f"{current[k]:.3f}→{new_weights[k]:.3f}" for k in WEIGHT_KEYS if abs(new_weights[k] - current[k]) > 0.001},
                )

        except Exception:
            logger.exception("Brain auto-tuner error")

        await asyncio.sleep(config.autotuner_interval)
