"""
Shared MLX LoRA training session: lock, status file, stdout progress + ETA, cancel.

Used by fine_tuner (chat-history LoRA) and tool_adapter_trainer (bundled tool-calling corpus).
"""
from __future__ import annotations

import json
import os
import re
import subprocess
import time
from dataclasses import dataclass, field
from threading import Lock
from typing import Any

from app_mlx import MLX_DATA_DIR

STATUS_FILE = os.path.join(MLX_DATA_DIR, "fine_tune_status.json")

_training_lock = Lock()
_training_process: subprocess.Popen | None = None
_training_cancelled = False
_session_active = False

_MAX_SUBPROC_LOG_LINES = 60
_MAX_SUBPROC_LINE_CHARS = 2000
_MAX_STATUS_MESSAGE_CHARS = 4500
_STALE_STATUS_SECONDS = 120
# Preflight/extracting may run for minutes before mlx_lm.lora exists (stop server, copy data).
_STALE_PREP_NO_PROC_SECONDS = 600
_ETA_WARMUP_ITERS = 5
_ETA_EMA_ALPHA = 0.25
_ETA_RECENT_WINDOW = 12
_ETA_MAX_SEC_PER_ITER = 90.0

_ITER_LOSS_RE = re.compile(r"Iter\s+(\d+).*?Train\s+loss\s+([\d.]+)", re.IGNORECASE)
_ITER_FRAC_RE = re.compile(r"(\d+)\s*/\s*(\d+)")
_IT_PER_SEC_RE = re.compile(r"It/sec\s+([\d.]+)", re.IGNORECASE)


@dataclass
class LoraParseState:
    """Mutable parser state for ETA across stdout lines."""

    ema_seconds_per_iter: float | None = None
    recent_seconds_per_iter: list[float] = field(default_factory=list)
    last_iter_monotonic: float | None = None
    last_iteration: int = 0


def format_subprocess_tail(
    lines: list[str],
    max_lines: int = _MAX_SUBPROC_LOG_LINES,
    max_chars: int = 3500,
) -> str:
    if not lines:
        return ""
    tail = lines[-max_lines:]
    s = "\n".join(tail)
    if len(s) > max_chars:
        s = s[-max_chars:]
    return s.strip()


def emit_training_activity(data: dict) -> None:
    try:
        from activity_bus import emit_activity_enriched

        kind = str(data.get("kind") or "chat").strip().lower()
        phase = str(data.get("phase") or "").strip().lower()
        adapter_name = str(data.get("adapter_name") or "mlx")
        label = "tool adapter" if kind == "tool_adapter" else "fine-tuning"
        summary = f"MLX {label} update"
        state = "progress"
        if phase == "preflight":
            state = "started"
            summary = f"MLX {label} preparing ({adapter_name})"
        elif phase == "extracting":
            state = "started"
            summary = f"MLX {label} started ({adapter_name})"
        elif phase == "training":
            iteration = int(data.get("iteration") or 0)
            total_iters = int(data.get("total_iters") or 0)
            if total_iters > 0:
                summary = f"MLX {label} {iteration}/{total_iters}"
            else:
                summary = f"MLX {label} iter {iteration}"
        elif phase == "completed":
            state = "completed"
            summary = f"MLX {label} completed ({adapter_name})"
        elif phase == "cancelled":
            state = "stopped"
            summary = f"MLX {label} cancelled ({adapter_name})"
        elif phase == "error":
            state = "error"
            summary = f"MLX {label} failed ({adapter_name})"
            err_msg = (data.get("message") or "").strip().replace("\n", " ")
            if err_msg:
                if len(err_msg) > 180:
                    err_msg = err_msg[:177] + "…"
                summary = f"{summary}: {err_msg}"
        emit_activity_enriched(
            {
                "source": "mlx",
                "state": state,
                "summary": summary,
                "job_id": adapter_name,
                "progress": data.get("progress"),
                "meta": {
                    "kind": kind,
                    "phase": phase,
                    "iteration": data.get("iteration"),
                    "total_iters": data.get("total_iters"),
                    "loss": data.get("loss"),
                    "message": data.get("message"),
                },
            },
            tags=["background"],
        )
    except Exception:
        pass


def write_status(data: dict) -> bool:
    os.makedirs(MLX_DATA_DIR, exist_ok=True)
    try:
        with open(STATUS_FILE, "w", encoding="utf-8") as f:
            json.dump(data, f)
    except Exception:
        return False
    emit_training_activity(data)
    return True


def _process_alive() -> bool:
    proc = _training_process
    if proc is None:
        return False
    return proc.poll() is None


def end_training_session() -> None:
    global _session_active
    with _training_lock:
        _session_active = False


def reconcile_stale_status(status: dict) -> dict:
    """If status claims running but subprocess is gone, mark interrupted."""
    if not status.get("running"):
        return status
    if _process_alive():
        return status
    phase = str(status.get("phase") or "").lower()
    try:
        age = time.time() - os.path.getmtime(STATUS_FILE)
    except OSError:
        age = _STALE_PREP_NO_PROC_SECONDS + 1
    # Legitimate preflight/extracting: no subprocess until mlx_lm.lora starts.
    if phase in ("preflight", "extracting") and age < _STALE_PREP_NO_PROC_SECONDS:
        return status
    # Training phase (or stale preflight past grace): no process means interrupted.
    out = dict(status)
    out["running"] = False
    out["phase"] = "error"
    out["message"] = "Training interrupted (process no longer running)."
    write_status(out)
    end_training_session()
    return out


def get_training_status(*, reconcile: bool = True) -> dict:
    try:
        if os.path.isfile(STATUS_FILE):
            with open(STATUS_FILE, encoding="utf-8") as f:
                status = json.load(f)
            if reconcile:
                return reconcile_stale_status(status)
            return status
    except Exception:
        pass
    return {"running": False}


def training_in_progress() -> bool:
    return bool(get_training_status(reconcile=True).get("running"))


def _training_busy_response(kind: str) -> dict:
    return {
        "ok": False,
        "message": f"Training already in progress ({kind}).",
        "busy_kind": kind,
        "status_code": 409,
    }


def try_acquire_training_slot() -> dict | None:
    """Return error dict if busy; None if slot reserved (caller must write status or end session)."""
    global _training_cancelled
    status = get_training_status(reconcile=True)
    with _training_lock:
        if _session_active or status.get("running"):
            kind = str(status.get("kind") or "chat")
            return _training_busy_response(kind)
        _training_cancelled = False
    return None


def begin_training_session(initial_status: dict) -> dict | None:
    """Atomically reserve the training slot and write initial status. Returns busy dict or None."""
    global _training_cancelled, _session_active
    status = get_training_status(reconcile=True)
    with _training_lock:
        if _session_active:
            kind = str(initial_status.get("kind") or "chat")
            return _training_busy_response(kind)
        if status.get("running"):
            kind = str(status.get("kind") or "chat")
            return _training_busy_response(kind)
        _training_cancelled = False
    if not write_status(initial_status):
        return {
            "ok": False,
            "message": "Could not write training status file.",
            "status_code": 500,
        }
    with _training_lock:
        _session_active = True
    return None


def _restart_mlx_after_tool_adapter_cancel() -> None:
    """Tool-adapter training stops MLX up front; restore server after cancel."""
    try:
        from tool_adapter_trainer import restart_mlx_server_after_training_cancel

        ok, msg = restart_mlx_server_after_training_cancel()
        if ok:
            return
        cur = get_training_status(reconcile=False)
        if str(cur.get("phase") or "").lower() != "cancelled":
            return
        cur["message"] = f"MLX server failed to restart: {msg}"
        write_status(cur)
    except Exception:
        pass


def cancel_training() -> dict:
    global _training_cancelled
    status = get_training_status(reconcile=False)
    kind = str(status.get("kind") or "chat")
    adapter_name = str(status.get("adapter_name") or "")
    adapter_path = str(status.get("adapter_path") or "")
    with _training_lock:
        _training_cancelled = True
        proc = _training_process
        if proc is not None and proc.poll() is None:
            try:
                proc.terminate()
            except Exception:
                pass
            return {"ok": True, "message": "Cancellation requested."}
    if status.get("running"):
        write_status({
            "running": False,
            "phase": "cancelled",
            "kind": kind,
            "adapter_name": adapter_name,
            "adapter_path": adapter_path,
        })
        end_training_session()
        if kind == "tool_adapter":
            _restart_mlx_after_tool_adapter_cancel()
        return {"ok": True, "message": "Cancellation requested."}
    end_training_session()
    return {"ok": False, "message": "No training in progress."}


def is_training_cancelled() -> bool:
    return _training_cancelled


def _eta_from_it_per_sec(it_per_sec: float, iteration: int, total_iters: int) -> int | None:
    if it_per_sec <= 0 or total_iters <= iteration:
        return None
    return max(0, int(round((total_iters - iteration) / it_per_sec)))


def _median_seconds_per_iter(samples: list[float]) -> float | None:
    if not samples:
        return None
    ordered = sorted(samples)
    mid = len(ordered) // 2
    if len(ordered) % 2 == 1:
        return ordered[mid]
    return (ordered[mid - 1] + ordered[mid]) / 2.0


def update_eta(
    state: LoraParseState,
    iteration: int,
    total_iters: int,
) -> int | None:
    now = time.monotonic()
    if state.last_iter_monotonic is not None and iteration > state.last_iteration:
        iter_delta = iteration - state.last_iteration
        dt = now - state.last_iter_monotonic
        if dt > 0 and iter_delta > 0:
            per_iter = dt / iter_delta
            if 0 < per_iter <= _ETA_MAX_SEC_PER_ITER:
                state.recent_seconds_per_iter.append(per_iter)
                if len(state.recent_seconds_per_iter) > _ETA_RECENT_WINDOW:
                    state.recent_seconds_per_iter.pop(0)
                if state.ema_seconds_per_iter is None:
                    state.ema_seconds_per_iter = per_iter
                else:
                    a = _ETA_EMA_ALPHA
                    state.ema_seconds_per_iter = a * per_iter + (1 - a) * state.ema_seconds_per_iter
    state.last_iter_monotonic = now
    state.last_iteration = iteration
    if iteration < _ETA_WARMUP_ITERS or total_iters <= iteration:
        return None
    rate = _median_seconds_per_iter(state.recent_seconds_per_iter)
    if rate is None:
        rate = state.ema_seconds_per_iter
    if rate is None or rate <= 0:
        return None
    return max(0, int(round((total_iters - iteration) * rate)))


def _parse_it_per_sec(line: str) -> float | None:
    m = _IT_PER_SEC_RE.search(line)
    if not m:
        return None
    try:
        val = float(m.group(1))
    except ValueError:
        return None
    return val if val > 0 else None


def _resolve_eta_seconds(
    *,
    line: str,
    state: LoraParseState,
    iteration: int,
    total_iters: int,
) -> int | None:
    it_per_sec = _parse_it_per_sec(line)
    if it_per_sec is not None:
        eta = _eta_from_it_per_sec(it_per_sec, iteration, total_iters)
        if eta is not None:
            return eta
    return update_eta(state, iteration, total_iters)


def parse_lora_line(
    line: str,
    *,
    state: LoraParseState,
    total_iters: int,
) -> dict[str, Any] | None:
    line_stripped = line.strip()
    if not line_stripped:
        return None
    m = _ITER_LOSS_RE.search(line_stripped)
    if m:
        iteration = min(int(m.group(1)), total_iters) if total_iters > 0 else int(m.group(1))
        loss = float(m.group(2))
        progress = min(1.0, iteration / total_iters) if total_iters > 0 else 0.0
        eta_seconds = _resolve_eta_seconds(
            line=line_stripped,
            state=state,
            iteration=iteration,
            total_iters=total_iters,
        )
        return {
            "iteration": iteration,
            "total_iters": total_iters,
            "loss": round(loss, 6),
            "progress": round(progress, 4),
            "eta_seconds": eta_seconds,
        }
    m2 = _ITER_FRAC_RE.search(line_stripped)
    if m2 and total_iters <= 0:
        iteration = int(m2.group(1))
        total_iters = int(m2.group(2))
        progress = min(1.0, iteration / total_iters) if total_iters > 0 else 0.0
        eta_seconds = _resolve_eta_seconds(
            line=line_stripped,
            state=state,
            iteration=iteration,
            total_iters=total_iters,
        )
        return {
            "iteration": iteration,
            "total_iters": total_iters,
            "loss": None,
            "progress": round(progress, 4),
            "eta_seconds": eta_seconds,
        }
    return None


def stream_lora_output(
    proc: subprocess.Popen,
    *,
    kind: str,
    total_iters: int,
    adapter_name: str,
    extra_status: dict | None = None,
) -> list[str]:
    """Read stdout until process ends; update status file. Returns log lines."""
    state = LoraParseState()
    captured: list[str] = []
    base_extra = dict(extra_status or {})
    if proc.stdout is None:
        return captured
    for line in proc.stdout:
        if is_training_cancelled():
            break
        raw = line.rstrip("\n\r")
        if raw.strip():
            captured.append(raw.strip()[:_MAX_SUBPROC_LINE_CHARS])
        parsed = parse_lora_line(raw, state=state, total_iters=total_iters)
        if parsed:
            payload = {
                "running": True,
                "phase": "training",
                "kind": kind,
                "adapter_name": adapter_name,
                **base_extra,
                **parsed,
            }
            write_status(payload)
    return captured


def run_lora_subprocess(
    cmd: list[str],
    *,
    kind: str,
    total_iters: int,
    adapter_name: str,
    cwd: str | None = None,
    env: dict[str, str] | None = None,
    extra_status: dict | None = None,
) -> tuple[int, list[str]]:
    """
    Run mlx_lm.lora under the global training lock. Caller must have acquired slot via try_acquire_training_slot.

    Returns (exit_code, log_lines).
    """
    global _training_process

    merged_env = None
    if env:
        merged_env = os.environ.copy()
        merged_env.update(env)

    with _training_lock:
        if _training_cancelled:
            return -1, []
        _training_process = subprocess.Popen(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            cwd=cwd or MLX_DATA_DIR,
            env=merged_env,
        )
        proc = _training_process

    log_lines = stream_lora_output(
        proc,
        kind=kind,
        total_iters=total_iters,
        adapter_name=adapter_name,
        extra_status=extra_status,
    )

    if is_training_cancelled() and proc.poll() is None:
        try:
            proc.terminate()
            proc.wait(timeout=5)
        except Exception:
            try:
                proc.kill()
            except Exception:
                pass

    try:
        exit_code = proc.wait(timeout=30)
    except Exception:
        exit_code = proc.poll() if proc.poll() is not None else -1

    with _training_lock:
        _training_process = None

    return exit_code, log_lines


def finish_training_status(
    *,
    kind: str,
    adapter_name: str,
    adapter_path: str,
    exit_code: int,
    log_lines: list[str],
    cancelled: bool | None = None,
) -> None:
    if cancelled is None:
        cancelled = _training_cancelled
    if cancelled:
        write_status({
            "running": False,
            "phase": "cancelled",
            "kind": kind,
            "adapter_name": adapter_name,
            "adapter_path": adapter_path,
        })
        if kind == "tool_adapter":
            _restart_mlx_after_tool_adapter_cancel()
    elif exit_code == 0:
        write_status({
            "running": False,
            "phase": "completed",
            "kind": kind,
            "adapter_name": adapter_name,
            "adapter_path": adapter_path,
        })
    else:
        tail = format_subprocess_tail(log_lines)
        msg = f"Training exited with code {exit_code}"
        if tail:
            msg = f"{msg}.\n--- mlx_lm.lora output (tail) ---\n{tail}"
        if len(msg) > _MAX_STATUS_MESSAGE_CHARS:
            msg = msg[: _MAX_STATUS_MESSAGE_CHARS - 1] + "…"
        write_status({
            "running": False,
            "phase": "error",
            "kind": kind,
            "adapter_name": adapter_name,
            "message": msg,
        })
    end_training_session()
