"""Hardware-agnostic pieces of the fine-tuning pipeline, shared by every backend
(app/services/finetune/mlx_backend.py, cuda_backend.py): exporting training data,
downloading base models, converting a fused model to GGUF, and importing it into
Ollama. Only the actual training/fusing step (run_lora_training/run_fuse) differs per
backend — see app/services/finetune/__init__.py for how a backend is selected.
"""

import json
import subprocess
import uuid
from pathlib import Path
from typing import Callable

from sqlalchemy.orm import Session

from app.core.config import get_settings
from app.models.dataset import DatasetExample, ExampleStatus

settings = get_settings()

MIN_APPROVED_EXAMPLES = 6


class FineTuneError(RuntimeError):
    pass


def job_dir(job_id: uuid.UUID) -> Path:
    d = Path(settings.finetune_data_dir) / str(job_id)
    d.mkdir(parents=True, exist_ok=True)
    return d


def export_training_data(dataset_id: uuid.UUID, out_dir: Path, db: Session) -> tuple[int, int]:
    examples = (
        db.query(DatasetExample)
        .filter(DatasetExample.dataset_id == dataset_id, DatasetExample.status == ExampleStatus.approved)
        .order_by(DatasetExample.created_at)
        .all()
    )
    if len(examples) < MIN_APPROVED_EXAMPLES:
        raise FineTuneError(
            f"Need at least {MIN_APPROVED_EXAMPLES} approved examples to fine-tune "
            f"(found {len(examples)}). Approve more examples in Dataset Manager first."
        )

    # Every backend's trainer requires each split to have at least `batch_size`
    # examples; keep valid at least 2 (the smallest batch size the UI allows) whenever
    # there's enough data, without starving train below MIN_APPROVED_EXAMPLES - 2.
    split = max(2, round(len(examples) * 0.15))
    split = min(split, len(examples) - 2)
    valid_examples = examples[:split]
    train_examples = examples[split:]

    def _write(path: Path, rows: list[DatasetExample]) -> None:
        with path.open("w") as f:
            for ex in rows:
                user_content = ex.instruction if not ex.input else f"{ex.instruction}\n\n{ex.input}"
                record = {
                    "messages": [
                        {"role": "user", "content": user_content},
                        {"role": "assistant", "content": ex.output},
                    ]
                }
                f.write(json.dumps(record, ensure_ascii=False) + "\n")

    _write(out_dir / "train.jsonl", train_examples)
    _write(out_dir / "valid.jsonl", valid_examples)
    return len(train_examples), len(valid_examples)


def ensure_base_model_downloaded(base_model: str) -> None:
    # Deliberately a lazy import: huggingface_hub is a hardware-specific extra (see
    # requirements-macos.txt / requirements-cuda.txt) since it's only needed for this
    # fine-tuning step. A top-level import would crash the entire app on boot anywhere
    # that dependency isn't installed (e.g. the core Docker image), not just this one
    # feature — found the hard way when the Docker image first booted.
    try:
        from huggingface_hub import snapshot_download
    except ImportError as exc:
        raise FineTuneError(
            "Fine-tuning isn't available on this install — huggingface_hub isn't "
            "installed. See README 'Fine-Tuning setup'."
        ) from exc

    # The training step's own download during training only fetches files needed to
    # load the model; the later fuse step (for the MLX backend) checks the snapshot
    # against the full repo tree with local_files_only=True and fails on missing
    # incidental files (README, .gitattributes) unless we've pulled the complete
    # snapshot up front. Harmless no-op cost for backends that don't need this.
    snapshot_download(base_model)


def patch_tokenizer_class(fused_dir: Path) -> None:
    """Whatever wrote tokenizer_config.json in the fused model dir stamps it with
    *its own* transformers version's tokenizer_class (e.g. "TokenizersBackend" on
    transformers 5.x), which the older transformers pinned in the isolated
    GGUF-conversion venv doesn't recognize. Force a generic, version-stable class name
    — the actual tokenizer behavior comes from tokenizer.json, not this field. Applies
    to every backend's fuse step, not just MLX's."""
    config_path = fused_dir / "tokenizer_config.json"
    if not config_path.exists():
        return
    data = json.loads(config_path.read_text())
    if "tokenizer_class" in data:
        data["tokenizer_class"] = "PreTrainedTokenizerFast"
        config_path.write_text(json.dumps(data, indent=2))


def run_gguf_convert(fused_dir: Path, gguf_path: Path) -> None:
    # Must stay a symlink path, not the fully-resolved realpath: the venv's
    # bin/python3 -> python3.14 -> /opt/homebrew/.../python3.14 symlink chain is how
    # Python finds this venv's pyvenv.cfg and site-packages (torch, gguf, etc.).
    # Resolving through to the real Homebrew binary makes it start up as a bare system
    # interpreter with none of the isolated venv's packages installed.
    python_bin = Path(settings.gguf_convert_python).absolute()
    script = Path(settings.gguf_convert_script).resolve()
    if not python_bin.exists() or not script.exists():
        raise FineTuneError(
            "GGUF conversion isn't set up on this machine. Run the one-time setup in "
            "the README ('Fine-Tuning setup') to create backend/.venv-gguf and clone "
            "the llama.cpp conversion script."
        )
    cmd = [str(python_bin), str(script), str(fused_dir), "--outfile", str(gguf_path), "--outtype", "f16"]
    run_streaming(cmd, lambda _line: None)
    if not gguf_path.exists():
        raise FineTuneError("GGUF conversion did not produce an output file.")


def import_into_ollama(tag: str, gguf_path: Path, system_prompt: str | None = None) -> None:
    # Ollama's HTTP /api/create only accepts a FROM referencing an already-pulled model
    # or an uploaded blob digest — a raw local file path needs the blob-upload dance the
    # `ollama` CLI already does internally, so shell out to it instead of reimplementing
    # that (confirmed working: `FROM <path>` resolves fine via the CLI).
    modelfile_lines = [f"FROM {gguf_path.resolve()}"]
    if system_prompt:
        escaped = system_prompt.replace('"""', '\\"\\"\\"')
        modelfile_lines.append(f'SYSTEM """{escaped}"""')
    modelfile_path = gguf_path.parent / "Modelfile"
    modelfile_path.write_text("\n".join(modelfile_lines))

    result = subprocess.run(
        ["ollama", "create", tag, "-f", str(modelfile_path)],
        capture_output=True,
        text=True,
        timeout=300,
    )
    if result.returncode != 0:
        raise FineTuneError(f"ollama create failed: {result.stderr or result.stdout}")


def run_streaming(cmd: list[str], on_line: Callable[[str], None]) -> None:
    proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1)
    assert proc.stdout is not None
    tail: list[str] = []
    for raw_line in proc.stdout:
        line = raw_line.rstrip("\n")
        if line:
            on_line(line)
            tail.append(line)
            if len(tail) > 25:
                tail.pop(0)
    returncode = proc.wait()
    if returncode != 0:
        detail = "\n".join(tail[-8:]) or "(no output captured)"
        raise FineTuneError(f"Command failed (exit {returncode}): {' '.join(cmd)}\n{detail}")