"""
Fine-tuning pipeline for MLX-based LoRA personalization.

Extracts training data from chat_storage (encrypted SQLite), writes
ephemeral JSONL files to a secure temp directory, runs mlx_lm.lora
as a subprocess in the MLX venv, and cleans up immediately.

Privacy / secure temp lifecycle:
  - Training data written to ~/.mugi/mlx/.training_tmp/ (mode 0o700 / 0o600)
  - Deleted in a finally block after training completes (success or failure)
  - atexit + signal handlers registered as safety net
  - sweep_stale_training_data() runs on backend startup to clean up after crashes
"""
from __future__ import annotations

import atexit
import json
import os
import signal
from threading import Thread
from typing import Any

from app_mlx import MLX_ADAPTERS_DIR, MLX_DATA_DIR, MLX_TRAINING_TMP_DIR, MLX_VENV_DIR
from mlx_training_run import (
    begin_training_session,
    cancel_training,
    end_training_session,
    finish_training_status,
    get_training_status,
    is_training_cancelled,
    run_lora_subprocess,
    write_status,
)
from utils.transcript_roles import is_real_user_message

# ---------------------------------------------------------------------------
# Secure temp directory
# ---------------------------------------------------------------------------

_TRAINING_DIR_MODE = 0o700
_TRAINING_FILE_MODE = 0o600


def _ensure_training_tmp() -> str:
    """Create the secure training temp directory. Returns its path."""
    os.makedirs(MLX_TRAINING_TMP_DIR, mode=_TRAINING_DIR_MODE, exist_ok=True)
    try:
        os.chmod(MLX_TRAINING_TMP_DIR, _TRAINING_DIR_MODE)
    except OSError:
        pass
    return MLX_TRAINING_TMP_DIR


def _cleanup_training_tmp():
    """Delete all contents of the training temp directory."""
    if not os.path.isdir(MLX_TRAINING_TMP_DIR):
        return
    import shutil
    try:
        shutil.rmtree(MLX_TRAINING_TMP_DIR, ignore_errors=True)
    except Exception:
        pass


def sweep_stale_training_data():
    """Remove any leftover training data from crashes. Call on backend startup."""
    _cleanup_training_tmp()


# Register safety nets
atexit.register(_cleanup_training_tmp)
for _sig in (signal.SIGTERM, signal.SIGINT):
    try:
        _prev = signal.getsignal(_sig)

        def _handler(signum, frame, prev=_prev):
            _cleanup_training_tmp()
            if callable(prev) and prev not in (signal.SIG_IGN, signal.SIG_DFL):
                prev(signum, frame)

        signal.signal(_sig, _handler)
    except (OSError, ValueError):
        pass


def get_fine_tune_status() -> dict:
    """Return the current fine-tuning status dict."""
    return get_training_status(reconcile=True)


def cancel_fine_tune() -> dict:
    """Cancel an in-progress fine-tuning run."""
    return cancel_training()


# ---------------------------------------------------------------------------
# Data extraction
# ---------------------------------------------------------------------------


def _extract_text(content: Any) -> str:
    """Extract plain text from message content (str or multipart list)."""
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        parts = []
        for block in content:
            if isinstance(block, dict):
                if block.get("type") == "text":
                    parts.append(block.get("text", ""))
                elif block.get("type") != "image_url":
                    parts.append(block.get("text", str(block)))
            else:
                parts.append(str(block))
        return " ".join(parts)
    return str(content) if content else ""


def _is_real_user_query_for_training(m: dict) -> bool:
    """True if this is a human user turn, not tool output stored under user role."""
    return is_real_user_message(m)


def _normalize_conversation_turns(turns: list[dict]) -> list[dict] | None:
    """Drop leading assistant-only prefix; require user-first dialog with ≥1 assistant."""
    while turns and turns[0].get("role") == "assistant":
        turns = turns[1:]
    if not turns or turns[0].get("role") != "user":
        return None
    if not any(t.get("role") == "assistant" for t in turns):
        return None
    if len(turns) < 2:
        return None
    return turns


def extract_training_data(
    thread_ids: list[str] | None = None,
    output_dir: str | None = None,
    validation_split: float = 0.1,
) -> tuple[str, int]:
    """Extract chat history into MLX-LM chat JSONL format."""
    from chat_storage import get_messages, _get_connection

    if output_dir is None:
        output_dir = _ensure_training_tmp()
    else:
        os.makedirs(output_dir, mode=_TRAINING_DIR_MODE, exist_ok=True)

    if thread_ids is None:
        conn = _get_connection()
        try:
            rows = conn.execute("SELECT DISTINCT thread_id FROM messages").fetchall()
            thread_ids = [r[0] for r in rows]
        finally:
            conn.close()

    conversations: list[list[dict]] = []
    for tid in thread_ids:
        msgs = get_messages(tid)
        turns: list[dict] = []
        for m in msgs:
            role = m.get("role", "").strip()
            if role not in ("user", "assistant"):
                continue
            if role == "user" and not _is_real_user_query_for_training(m):
                continue
            text = _extract_text(m.get("content", "")).strip()
            if not text or text.lower() == "init":
                continue
            turns.append({"role": role, "content": text})
        normalized = _normalize_conversation_turns(turns)
        if normalized is not None:
            conversations.append(normalized)

    if not conversations:
        return output_dir, 0

    split_idx = max(1, int(len(conversations) * (1 - validation_split)))
    train_convos = conversations[:split_idx]
    valid_convos = conversations[split_idx:] if len(conversations) > 1 else []

    train_path = os.path.join(output_dir, "train.jsonl")
    _write_jsonl(train_path, train_convos)

    if valid_convos:
        valid_path = os.path.join(output_dir, "valid.jsonl")
        _write_jsonl(valid_path, valid_convos)

    return output_dir, len(conversations)


def _write_jsonl(path: str, conversations: list[list[dict]]):
    """Write conversations to JSONL file in MLX-LM chat format."""
    fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, _TRAINING_FILE_MODE)
    try:
        with os.fdopen(fd, "w") as f:
            for convo in conversations:
                record = {"messages": convo}
                f.write(json.dumps(record) + "\n")
    except Exception:
        try:
            os.close(fd)
        except Exception:
            pass
        raise


# ---------------------------------------------------------------------------
# Training pipeline
# ---------------------------------------------------------------------------


def run_fine_tune(
    model_path: str,
    adapter_name: str,
    thread_ids: list[str] | None = None,
    iters: int = 200,
    learning_rate: float = 1e-5,
    num_layers: int = 4,
    batch_size: int = 1,
) -> dict:
    """Start a fine-tuning run in a background thread."""
    if not os.path.isfile(os.path.join(MLX_VENV_DIR, "bin", "python")):
        return {"ok": False, "message": "MLX venv not installed."}

    if not adapter_name or ".." in adapter_name or "/" in adapter_name:
        return {"ok": False, "message": "Invalid adapter name."}

    adapter_output = os.path.join(MLX_ADAPTERS_DIR, adapter_name)

    busy = begin_training_session({
        "running": True,
        "phase": "extracting",
        "kind": "chat",
        "iteration": 0,
        "total_iters": iters,
        "loss": None,
        "eta_seconds": None,
        "adapter_name": adapter_name,
    })
    if busy:
        return busy

    def _train():
        data_dir = None
        try:
            data_dir, total = extract_training_data(thread_ids=thread_ids)
            if is_training_cancelled():
                finish_training_status(
                    kind="chat",
                    adapter_name=adapter_name,
                    adapter_path=adapter_output,
                    exit_code=-1,
                    log_lines=[],
                )
                return
            if total == 0:
                write_status({
                    "running": False,
                    "phase": "error",
                    "kind": "chat",
                    "adapter_name": adapter_name,
                    "message": "No training data found.",
                })
                end_training_session()
                return

            if is_training_cancelled():
                finish_training_status(
                    kind="chat",
                    adapter_name=adapter_name,
                    adapter_path=adapter_output,
                    exit_code=-1,
                    log_lines=[],
                )
                return

            write_status({
                "running": True,
                "phase": "training",
                "kind": "chat",
                "iteration": 0,
                "total_iters": iters,
                "loss": None,
                "eta_seconds": None,
                "adapter_name": adapter_name,
                "conversations": total,
            })

            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", data_dir,
                "--adapter-path", adapter_output,
                "--iters", str(iters),
                "--learning-rate", str(learning_rate),
                "--num-layers", str(num_layers),
                "--batch-size", str(batch_size),
                "--mask-prompt",
            ]

            exit_code, log_lines = run_lora_subprocess(
                cmd,
                kind="chat",
                total_iters=iters,
                adapter_name=adapter_name,
                cwd=MLX_DATA_DIR,
                extra_status={"conversations": total},
            )

            finish_training_status(
                kind="chat",
                adapter_name=adapter_name,
                adapter_path=adapter_output,
                exit_code=exit_code,
                log_lines=log_lines,
            )
        except Exception as e:
            write_status({
                "running": False,
                "phase": "error",
                "kind": "chat",
                "adapter_name": adapter_name,
                "message": str(e)[:500],
            })
            end_training_session()
        finally:
            _cleanup_training_tmp()

    Thread(target=_train, daemon=True).start()
    return {"ok": True, "message": "Fine-tuning started."}
