"""
GRPO training runner for the evolution pipeline.

Builds a GRPO dataset from recent trajectories + rewards, then runs
MLX-native GRPO training as a subprocess. Output is a LoRA adapter
saved per base model under ~/.mugi/mlx/adapters/.

Pattern mirrors fine_tuner.py: subprocess, status file, lock, cleanup.
"""
from __future__ import annotations

import atexit
import json
import os
import re
import shutil
import subprocess
import time
from datetime import datetime
from threading import Lock

from evolution.meta import (
    MLX_ADAPTERS_DIR,
    _EVOLUTION_DIR,
    model_slug,
    read_evolution_meta,
    write_evolution_meta,
)
from evolution.rewards import compute_reward
from evolution.trajectory_logger import TRAJECTORIES_DIR, load_trajectory_records

_GRPO_DATA_DIR = os.path.join(_EVOLUTION_DIR, "grpo_data")
_GRPO_STATUS_PATH = os.path.join(_EVOLUTION_DIR, "grpo_status.json")
_DIR_MODE = 0o700
_FILE_MODE = 0o600

_grpo_lock = Lock()
_grpo_process: subprocess.Popen | None = None

# MLX venv path (same as app_mlx.py)
_MLX_VENV_DIR = os.path.join(os.path.expanduser("~"), ".mugi", "mlx-venv")

# Default training hyperparameters
GRPO_DEFAULT_ITERS = 100
GRPO_DEFAULT_LEARNING_RATE = 1e-5
GRPO_DEFAULT_BATCH_SIZE = 1
GRPO_DEFAULT_NUM_LAYERS = 4


def _log_evolution(msg: str, level: str = "WARNING") -> None:
    try:
        from tools._logging import _debug_log
        _debug_log(f"[evolution/grpo_runner] {msg}", level)
    except Exception:
        pass


def _write_status(data: dict) -> None:
    os.makedirs(_EVOLUTION_DIR, mode=_DIR_MODE, exist_ok=True)
    try:
        fd = os.open(_GRPO_STATUS_PATH, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, _FILE_MODE)
        try:
            with os.fdopen(fd, "w") as f:
                json.dump(data, f)
        except Exception:
            try:
                os.close(fd)
            except Exception:
                pass
    except Exception as e:
        _log_evolution(f"_write_status failed: {e}", "ERROR")


def get_grpo_status() -> dict:
    try:
        if os.path.isfile(_GRPO_STATUS_PATH):
            with open(_GRPO_STATUS_PATH, "r") as f:
                return json.load(f)
    except Exception as e:
        _log_evolution(f"get_grpo_status read failed: {e}", "ERROR")
    return {"running": False}


def _cleanup_grpo_data() -> None:
    if os.path.isdir(_GRPO_DATA_DIR):
        try:
            shutil.rmtree(_GRPO_DATA_DIR, ignore_errors=True)
        except Exception as e:
            _log_evolution(f"_cleanup_grpo_data failed: {e}", "ERROR")


atexit.register(_cleanup_grpo_data)


def _build_grpo_dataset(trajectories: list[dict]) -> str:
    """Build a GRPO dataset from trajectories + rewards.

    Writes train.jsonl to _GRPO_DATA_DIR. Each line:
    { "messages": [...], "reward": float }

    Returns the dataset directory path.
    """
    os.makedirs(_GRPO_DATA_DIR, mode=_DIR_MODE, exist_ok=True)
    train_path = os.path.join(_GRPO_DATA_DIR, "train.jsonl")
    fd = os.open(train_path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, _FILE_MODE)
    try:
        with os.fdopen(fd, "w") as f:
            for traj in trajectories:
                reward = compute_reward(traj)
                messages = traj.get("messages", [])
                record = {"messages": messages, "reward": reward}
                f.write(json.dumps(record, default=str) + "\n")
    except Exception:
        try:
            os.close(fd)
        except Exception:
            pass
        raise
    return _GRPO_DATA_DIR


def _get_current_mlx_model() -> str:
    """Read current MLX model from app prefs."""
    try:
        from utils.app_prefs_store import load_app_prefs_cached
        prefs = load_app_prefs_cached()
        return (prefs.get("mlx_model") or "").strip()
    except Exception:
        pass
    return ""


def _has_grpo_library() -> bool:
    """Check if the MLX venv has a GRPO-capable library installed."""
    python = os.path.join(_MLX_VENV_DIR, "bin", "python")
    if not os.path.isfile(python):
        return False
    for module in ("mlxsmith", "mlx_guided_grpo", "mlx_lm"):
        try:
            result = subprocess.run(
                [python, "-c", f"import {module}"],
                capture_output=True, timeout=10,
            )
            if result.returncode == 0:
                return True
        except Exception:
            continue
    return False


def run_grpo_training() -> dict:
    """Run one GRPO training cycle. Returns {"ok": bool, "message": str}.

    Blocking call; meant to be called from the scheduler thread.
    """
    global _grpo_process

    with _grpo_lock:
        status = get_grpo_status()
        if status.get("running"):
            return {"ok": False, "message": "GRPO training already in progress."}

    model_path = _get_current_mlx_model()
    if not model_path:
        return {"ok": False, "message": "No MLX model configured."}

    if not _has_grpo_library():
        return {"ok": False, "message": "No GRPO library found in MLX venv."}

    meta = read_evolution_meta()
    trajectories = load_trajectory_records(since_iso=meta.get("last_training_at"))
    if not trajectories:
        return {"ok": False, "message": "No new trajectories to train on."}

    slug = model_slug(model_path)
    adapter_name = f"evolution_{slug}"
    adapter_output = os.path.join(MLX_ADAPTERS_DIR, adapter_name)

    _log_evolution(
        f"Starting GRPO training: model={model_path}, trajectories={len(trajectories)}, adapter={adapter_name}",
        "INFO",
    )

    _write_status({
        "running": True, "phase": "building_dataset",
        "model": model_path, "adapter_name": adapter_name,
    })

    try:
        dataset_dir = _build_grpo_dataset(trajectories)

        _write_status({
            "running": True, "phase": "training",
            "model": model_path, "adapter_name": adapter_name,
            "trajectories": len(trajectories),
        })

        os.makedirs(adapter_output, exist_ok=True)

        python = os.path.join(_MLX_VENV_DIR, "bin", "python")

        cmd = [
            python, "-m", "mlx_lm.lora",
            "--train",
            "--model", model_path,
            "--data", dataset_dir,
            "--adapter-path", adapter_output,
            "--iters", str(GRPO_DEFAULT_ITERS),
            "--learning-rate", str(GRPO_DEFAULT_LEARNING_RATE),
            "--num-layers", str(GRPO_DEFAULT_NUM_LAYERS),
            "--batch-size", str(GRPO_DEFAULT_BATCH_SIZE),
        ]

        with _grpo_lock:
            _grpo_process = subprocess.Popen(
                cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                text=True,
                cwd=_EVOLUTION_DIR,
            )

        iter_re = re.compile(r"Iter\s+(\d+).*?Train\s+loss\s+([\d.]+)", re.IGNORECASE)
        for line in _grpo_process.stdout:
            m = iter_re.search(line.strip())
            if m:
                iteration = int(m.group(1))
                loss = float(m.group(2))
                _write_status({
                    "running": True, "phase": "training",
                    "model": model_path, "adapter_name": adapter_name,
                    "iteration": iteration, "loss": round(loss, 6),
                })

        exit_code = _grpo_process.wait(timeout=600)
        with _grpo_lock:
            _grpo_process = None

        if exit_code != 0:
            _log_evolution(f"Training exited with code {exit_code}", "ERROR")
            _write_status({
                "running": False, "phase": "error",
                "message": f"Training exited with code {exit_code}",
            })
            return {"ok": False, "message": f"Training exited with code {exit_code}"}

        meta["last_training_at"] = datetime.now().isoformat()
        meta["last_trajectory_count"] = len(trajectories)
        adapters_map = meta.get("adapters_by_model") or {}
        adapters_map[model_path] = adapter_name
        meta["adapters_by_model"] = adapters_map
        write_evolution_meta(meta)

        _log_evolution(
            f"GRPO training complete: adapter={adapter_name}, trajectories={len(trajectories)}",
            "INFO",
        )

        _write_status({
            "running": False, "phase": "completed",
            "adapter_name": adapter_name,
            "adapter_path": adapter_output,
        })

        return {"ok": True, "message": f"GRPO training complete. Adapter: {adapter_name}"}

    except Exception as e:
        with _grpo_lock:
            _grpo_process = None
        _log_evolution(f"GRPO training failed: {e}", "ERROR")
        _write_status({
            "running": False, "phase": "error",
            "message": str(e)[:500],
        })
        return {"ok": False, "message": f"GRPO training failed: {e}"}

    finally:
        _cleanup_grpo_data()
