"""
Train mugi-agent-tools-* LoRA adapters from the bundled synthetic tool-calling corpus.

Invoked from mugi_backend routes and optionally CLI (scripts/train_mugi_tool_adapter.sh).
"""
from __future__ import annotations

import json
import os
import shutil
import subprocess
import sys
from threading import Thread
from typing import Any

from app_mlx import (
    APP_PREFS_PATH,
    MLX_ADAPTERS_DIR,
    MLX_DATA_DIR,
    MLX_VENV_DIR,
    _mlx_model_slug,
    _resolve_mlx_model_path,
    start_mlx_server,
    stop_mlx_server,
)
from mlx_training_run import (
    begin_training_session,
    end_training_session,
    finish_training_status,
    get_training_status,
    is_training_cancelled,
    run_lora_subprocess,
    write_status,
)

_TOOL_ADAPTER_KIND = "tool_adapter"
_USER_DATASET_DIR = os.path.join(MLX_DATA_DIR, "tool_adapter_dataset")
_TRAINING_FILE_MODE = 0o600
_TRAINING_DIR_MODE = 0o700
_DISK_WARN_BYTES = 2 * 1024 * 1024 * 1024
_DISK_BLOCK_BYTES = 500 * 1024 * 1024
_LEARNING_RATE_DEFAULT = 1e-5
_BATCH_DEFAULT = 1


def _load_prefs() -> dict:
    if not os.path.isfile(APP_PREFS_PATH):
        return {}
    try:
        with open(APP_PREFS_PATH, encoding="utf-8") as fh:
            data = json.load(fh)
        return data if isinstance(data, dict) else {}
    except Exception:
        return {}


def _apply_active_adapter_pref(adapter_name: str) -> None:
    from utils.app_prefs_store import mutate_app_prefs

    def _mut(prefs: dict) -> None:
        prefs["mlx_active_adapter"] = adapter_name

    mutate_app_prefs(_mut)


def _default_params(model_path: str) -> dict[str, int]:
    low = (model_path or "").lower()
    if any(tok in low for tok in ("26b", "27b", "31b", "35b", "70b")):
        return {"iters": 1500, "num_layers": 6}
    if any(tok in low for tok in ("9b", "8b", "7b", "4b", "3b")):
        return {"iters": 1200, "num_layers": 8}
    return {"iters": 1000, "num_layers": 8}


def _is_moe_model(model_path: str) -> bool:
    low = (model_path or "").lower()
    return any(tok in low for tok in ("a4b", "a3b", "a2b", "moe", "mixtral", "deepseek-moe"))


def _sanitize_adapter_name(name: str) -> str | None:
    name = (name or "").strip()
    if not name or ".." in name or "/" in name or "\\" in name:
        return None
    return name


def _adapter_is_complete(adapter_path: str) -> bool:
    if not os.path.isdir(adapter_path):
        return False
    config_path = os.path.join(adapter_path, "adapter_config.json")
    weights_path = os.path.join(adapter_path, "adapters.safetensors")
    return os.path.isfile(config_path) or os.path.isfile(weights_path)


def _dataset_source_dirs() -> list[str]:
    roots: list[str] = []
    repo = (os.environ.get("MUGI_REPO_ROOT") or "").strip()
    if repo:
        roots.append(os.path.join(repo, "training", "mlx_mugi_agent_tools"))
    cwd = os.getcwd()
    roots.append(os.path.join(cwd, "training", "mlx_mugi_agent_tools"))
    try:
        here = os.path.dirname(os.path.abspath(__file__))
        roots.append(os.path.join(here, "training", "mlx_mugi_agent_tools"))
    except Exception:
        pass
    seen: set[str] = set()
    out: list[str] = []
    for p in roots:
        ap = os.path.abspath(p)
        if ap in seen:
            continue
        seen.add(ap)
        if os.path.isdir(ap):
            out.append(ap)
    return out


def _user_dataset_ready() -> str | None:
    train_path = os.path.join(_USER_DATASET_DIR, "train.jsonl")
    valid_path = os.path.join(_USER_DATASET_DIR, "valid.jsonl")
    if os.path.isfile(train_path) and os.path.isfile(valid_path):
        return _USER_DATASET_DIR
    return None


def ensure_tool_adapter_dataset() -> tuple[str | None, str]:
    """Return (data_dir, error_message). Copies bundled JSONL into ~/.mugi on first use."""
    existing = _user_dataset_ready()
    if existing:
        return existing, ""

    for src in _dataset_source_dirs():
        train_src = os.path.join(src, "train.jsonl")
        valid_src = os.path.join(src, "valid.jsonl")
        if not (os.path.isfile(train_src) and os.path.isfile(valid_src)):
            continue
        try:
            os.makedirs(_USER_DATASET_DIR, mode=_TRAINING_DIR_MODE, exist_ok=True)
            for name in ("train.jsonl", "valid.jsonl"):
                dst = os.path.join(_USER_DATASET_DIR, name)
                shutil.copy2(os.path.join(src, name), dst)
                os.chmod(dst, _TRAINING_FILE_MODE)
            return _USER_DATASET_DIR, ""
        except Exception as e:
            return None, f"Failed to install training dataset: {e}"

    return None, (
        "Tool-calling training dataset not found. "
        "In development, run: python training/mlx_mugi_agent_tools/generate_dataset.py"
    )


def _disk_preflight() -> dict[str, Any]:
    path = MLX_DATA_DIR
    try:
        usage = shutil.disk_usage(path)
        free_bytes = int(usage.free)
        disk_ok = free_bytes >= _DISK_BLOCK_BYTES
        warn = free_bytes < _DISK_WARN_BYTES
        return {
            "disk_ok": disk_ok,
            "free_bytes": free_bytes,
            "disk_warning": warn,
        }
    except Exception:
        return {"disk_ok": True, "free_bytes": None, "disk_warning": False}


def _grpo_busy() -> bool:
    try:
        from evolution.grpo_runner import get_grpo_status
        return bool(get_grpo_status().get("running"))
    except Exception:
        return False


def _metal_toolchain_warning() -> str | None:
    try:
        proc = subprocess.run(
            ["xcrun", "metal", "-version"],
            capture_output=True,
            timeout=5,
            check=False,
        )
        if proc.returncode == 0:
            return None
    except Exception:
        pass
    return (
        "Metal Toolchain may be missing. Install via Xcode Settings → Components "
        "or: xcodebuild -downloadComponent MetalToolchain"
    )


def _mlx_venv_ready() -> bool:
    return os.path.isfile(os.path.join(MLX_VENV_DIR, "bin", "python"))


def resolve_tool_adapter_plan(
    model_path: str | None = None,
    *,
    iters: int | None = None,
    num_layers: int | None = None,
    learning_rate: float | None = None,
    batch_size: int | None = None,
) -> dict[str, Any]:
    prefs = _load_prefs()
    model = (model_path or prefs.get("mlx_model") or "").strip()
    if not model:
        return {"ok": False, "error": "No MLX model selected. Choose a model in Settings → LLMs."}

    if not _mlx_venv_ready():
        return {"ok": False, "error": "MLX venv not installed. Install MLX from Settings → LLMs."}

    if _grpo_busy():
        return {
            "ok": False,
            "error": "Evolution GRPO training is running. Wait for it to finish before tool-adapter training.",
            "grpo_busy": True,
        }

    resolved = _resolve_mlx_model_path(model)
    expanded = os.path.expanduser(resolved)
    if os.path.isabs(expanded) and not os.path.isdir(expanded):
        return {"ok": False, "error": f"Local model path does not exist: {resolved}"}

    slug = _mlx_model_slug(resolved)
    adapter_name = _sanitize_adapter_name(f"mugi-agent-tools-{slug}")
    if not adapter_name:
        return {"ok": False, "error": "Invalid adapter name derived from model."}
    if adapter_name.startswith("evolution_"):
        return {"ok": False, "error": "Refusing to train into an evolution_* adapter name."}

    data_dir, data_err = ensure_tool_adapter_dataset()
    dataset_ok = data_dir is not None and not data_err
    adapter_path = os.path.join(MLX_ADAPTERS_DIR, adapter_name)
    adapter_exists = os.path.isdir(adapter_path)
    adapter_complete = _adapter_is_complete(adapter_path) if adapter_exists else False

    defaults = _default_params(resolved)
    iters_val = int(iters if iters is not None else defaults["iters"])
    num_layers_val = int(num_layers if num_layers is not None else defaults["num_layers"])
    lr_val = float(learning_rate if learning_rate is not None else _LEARNING_RATE_DEFAULT)
    batch_val = int(batch_size if batch_size is not None else _BATCH_DEFAULT)

    warnings: list[str] = []
    if _is_moe_model(resolved):
        warnings.append("MoE model: training uses MLX_ENABLE_TF32=0 (required for MoE LoRA on current MLX).")
    metal_warn = _metal_toolchain_warning()
    if metal_warn:
        warnings.append(metal_warn)
    if any(tok in resolved.lower() for tok in ("70b", "72b")):
        warnings.append("Large model: training may take many hours and require substantial RAM.")

    disk = _disk_preflight()
    if disk.get("disk_warning"):
        warnings.append("Low disk space on ~/.mugi/mlx — free at least 2 GB before training.")
    if not disk.get("disk_ok"):
        return {
            "ok": False,
            "error": "Not enough free disk space under ~/.mugi/mlx (need ~500 MB minimum).",
            **disk,
        }

    train_n = valid_n = 0
    if data_dir:
        try:
            with open(os.path.join(data_dir, "train.jsonl"), encoding="utf-8") as f:
                train_n = sum(1 for _ in f)
            with open(os.path.join(data_dir, "valid.jsonl"), encoding="utf-8") as f:
                valid_n = sum(1 for _ in f)
        except Exception:
            pass

    return {
        "ok": True,
        "model": resolved,
        "adapter_name": adapter_name,
        "adapter_path": adapter_path,
        "data_dir": data_dir or "",
        "dataset_ok": dataset_ok,
        "dataset_error": data_err,
        "iters": iters_val,
        "num_layers": num_layers_val,
        "learning_rate": lr_val,
        "batch_size": batch_val,
        "is_moe": _is_moe_model(resolved),
        "adapter_exists": adapter_exists,
        "adapter_complete": adapter_complete,
        "grpo_busy": False,
        "train_examples": train_n,
        "valid_examples": valid_n,
        "warnings": warnings,
        **disk,
    }


def _context_length_from_prefs(prefs: dict) -> int | None:
    val = prefs.get("llm_context_length")
    try:
        val = int(val)
    except (TypeError, ValueError):
        return None
    if 16384 <= val <= 1_048_576:
        return val
    return None


def _abort_if_cancelled(adapter_name: str, adapter_path: str) -> bool:
    if not is_training_cancelled():
        return False
    finish_training_status(
        kind=_TOOL_ADAPTER_KIND,
        adapter_name=adapter_name,
        adapter_path=adapter_path,
        exit_code=-1,
        log_lines=[],
    )
    return True


def restart_mlx_server_after_training_cancel() -> tuple[bool, str]:
    """Best-effort MLX restart after tool-adapter training was cancelled (server was stopped)."""
    return _restart_mlx_after_train(_load_prefs())


def _restart_mlx_after_train(prefs: dict) -> tuple[bool, str]:
    model_path = (prefs.get("mlx_model") or "").strip()
    if not model_path:
        return False, "No mlx_model in preferences."
    model_type = (prefs.get("mlx_model_type") or "auto").strip()
    ok, msg = start_mlx_server(model_path, model_type, _context_length_from_prefs(prefs))
    return ok, msg


def run_tool_adapter_train(
    *,
    model_path: str | None = None,
    overwrite: bool = False,
    apply_when_finished: bool = True,
    iters: int | None = None,
    num_layers: int | None = None,
    learning_rate: float | None = None,
    batch_size: int | None = None,
    blocking: bool = False,
) -> dict:
    plan = resolve_tool_adapter_plan(
        model_path,
        iters=iters,
        num_layers=num_layers,
        learning_rate=learning_rate,
        batch_size=batch_size,
    )
    if not plan.get("ok"):
        return {
            "ok": False,
            "message": plan.get("error") or "Plan validation failed.",
            "status_code": 400,
        }

    if _grpo_busy():
        return {"ok": False, "message": "Evolution GRPO training is running.", "status_code": 409}

    adapter_name = plan["adapter_name"]
    adapter_path = plan["adapter_path"]
    if plan.get("adapter_complete") and not overwrite:
        return {
            "ok": False,
            "message": f"Adapter '{adapter_name}' already exists. Pass overwrite=true to replace it.",
            "status_code": 400,
        }

    if not plan.get("dataset_ok"):
        return {
            "ok": False,
            "message": plan.get("dataset_error") or "Training dataset not available.",
            "status_code": 400,
        }

    busy = begin_training_session({
        "running": True,
        "phase": "preflight",
        "kind": _TOOL_ADAPTER_KIND,
        "iteration": 0,
        "total_iters": plan["iters"],
        "adapter_name": adapter_name,
        "adapter_path": adapter_path,
    })
    if busy:
        return busy

    def _train():
        restart_warning = ""
        try:
            stop_mlx_server()

            if _abort_if_cancelled(adapter_name, adapter_path):
                return

            if os.path.isdir(adapter_path) and (
                overwrite or not _adapter_is_complete(adapter_path)
            ):
                shutil.rmtree(adapter_path, ignore_errors=True)
            os.makedirs(adapter_path, exist_ok=True)

            write_status({
                "running": True,
                "phase": "training",
                "kind": _TOOL_ADAPTER_KIND,
                "iteration": 0,
                "total_iters": plan["iters"],
                "loss": None,
                "eta_seconds": None,
                "adapter_name": adapter_name,
            })

            if _abort_if_cancelled(adapter_name, adapter_path):
                return

            python = os.path.join(MLX_VENV_DIR, "bin", "python")
            cmd = [
                python, "-m", "mlx_lm.lora",
                "--train",
                "--model", plan["model"],
                "--data", plan["data_dir"],
                "--adapter-path", adapter_path,
                "--iters", str(plan["iters"]),
                "--batch-size", str(plan["batch_size"]),
                "--num-layers", str(plan["num_layers"]),
                "--learning-rate", str(plan["learning_rate"]),
                "--mask-prompt",
            ]
            env = None
            if plan.get("is_moe"):
                env = {"MLX_ENABLE_TF32": "0"}

            exit_code, log_lines = run_lora_subprocess(
                cmd,
                kind=_TOOL_ADAPTER_KIND,
                total_iters=plan["iters"],
                adapter_name=adapter_name,
                cwd=MLX_DATA_DIR,
                env=env,
            )

            if is_training_cancelled():
                finish_training_status(
                    kind=_TOOL_ADAPTER_KIND,
                    adapter_name=adapter_name,
                    adapter_path=adapter_path,
                    exit_code=exit_code,
                    log_lines=log_lines,
                )
                return

            if exit_code == 0 and not _adapter_is_complete(adapter_path):
                finish_training_status(
                    kind=_TOOL_ADAPTER_KIND,
                    adapter_name=adapter_name,
                    adapter_path=adapter_path,
                    exit_code=1,
                    log_lines=log_lines,
                )
                cur = get_training_status(reconcile=False)
                cur["message"] = (
                    "Training finished but adapter weights are missing or incomplete."
                )
                write_status(cur)
                return

            if exit_code == 0:
                prefs = _load_prefs()
                if apply_when_finished:
                    _apply_active_adapter_pref(adapter_name)
                ok_restart, restart_msg = _restart_mlx_after_train(prefs)
                if not ok_restart:
                    restart_warning = (
                        f"Adapter saved but MLX server failed to restart: {restart_msg}"
                    )
                finish_training_status(
                    kind=_TOOL_ADAPTER_KIND,
                    adapter_name=adapter_name,
                    adapter_path=adapter_path,
                    exit_code=0,
                    log_lines=log_lines,
                )
                if restart_warning:
                    cur = get_training_status(reconcile=False)
                    cur["message"] = restart_warning
                    write_status(cur)
            else:
                finish_training_status(
                    kind=_TOOL_ADAPTER_KIND,
                    adapter_name=adapter_name,
                    adapter_path=adapter_path,
                    exit_code=exit_code,
                    log_lines=log_lines,
                )
        except Exception as e:
            write_status({
                "running": False,
                "phase": "error",
                "kind": _TOOL_ADAPTER_KIND,
                "adapter_name": adapter_name,
                "message": str(e)[:500],
            })
            end_training_session()

    worker = Thread(target=_train, daemon=not blocking)
    worker.start()
    if blocking:
        worker.join()
        status = get_training_status(reconcile=False)
        phase = str(status.get("phase") or "")
        if phase == "completed":
            return {
                "ok": True,
                "message": "Tool-adapter training completed.",
                "adapter_name": adapter_name,
            }
        if phase == "cancelled":
            return {"ok": False, "message": "Training cancelled.", "adapter_name": adapter_name}
        return {
            "ok": False,
            "message": status.get("message") or f"Training ended with phase {phase!r}.",
            "adapter_name": adapter_name,
        }
    return {"ok": True, "message": "Tool-adapter training started.", "adapter_name": adapter_name}


def main(argv: list[str] | None = None) -> int:
    """CLI entry for scripts/train_mugi_tool_adapter.sh."""
    import argparse
    parser = argparse.ArgumentParser(description="Train Mugi tool-calling LoRA adapter")
    parser.add_argument("--model", default="")
    parser.add_argument("--overwrite", action="store_true")
    parser.add_argument("--no-apply", action="store_true")
    parser.add_argument("--iters", type=int, default=None)
    parser.add_argument("--num-layers", type=int, default=None)
    args = parser.parse_args(argv)
    result = run_tool_adapter_train(
        model_path=args.model or None,
        overwrite=args.overwrite,
        apply_when_finished=not args.no_apply,
        iters=args.iters,
        num_layers=args.num_layers,
        blocking=True,
    )
    if not result.get("ok"):
        print(result.get("message") or "Failed", file=sys.stderr)
        return 1
    print(result.get("message") or "Started")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
