"""Output bundle helpers shared by the GGUF conversion scripts."""

from __future__ import annotations

import logging
import shutil
from pathlib import Path

log = logging.getLogger("nvfp4gguf.bundle")

CONTEXT_HPARAMS = (
    "max_position_embeddings",
    "n_ctx",
    "n_positions",
    "max_length",
    "max_sequence_length",
    "model_max_length",
)


def is_model_context_key(name: str) -> bool:
    """The model's max context, not RoPE's original context length."""
    return name.endswith(".context_length") and ".rope." not in name


def require_context_length(length: int | None) -> int | None:
    if length is None:
        return None
    if length <= 0:
        raise SystemExit("--context-length must be a positive token count")
    return int(length)


def scalar_context_length(value) -> int | None:
    if isinstance(value, bool) or isinstance(value, (list, dict)):
        return None
    try:
        length = int(value)
    except (TypeError, ValueError):
        return None
    if length <= 0:
        return None
    return length


def _hf_context_key(hparams: dict) -> tuple[str | None, int | None]:
    """Key llama.cpp's find_hparam will see, and its scalar length if it has one."""
    for key in CONTEXT_HPARAMS:
        if key in hparams:
            return key, scalar_context_length(hparams[key])
    for nested_name in ("text_config", "llm_config", "language_config"):
        nested = hparams.get(nested_name)
        if not isinstance(nested, dict):
            continue
        for key in CONTEXT_HPARAMS:
            if key in nested:
                return key, scalar_context_length(nested[key])
    return None, None


def apply_hf_context_length(hparams: dict, length: int | None) -> int | None:
    """Return the context llama.cpp will write, after an optional override.

    llama.cpp reads the first usual context key on the top-level config. A
    list or a missing key is not a usable length. The override is written
    where that lookup will find it.
    """
    requested = require_context_length(length)
    found_key, found_value = _hf_context_key(hparams)
    if requested is not None:
        hparams[found_key or "max_position_embeddings"] = requested
        log.info("context length %d", requested)
        return requested
    if found_value is None:
        return None
    if found_key not in hparams:
        hparams[found_key] = found_value
        log.info("context length %d from the nested checkpoint config", found_value)
    else:
        log.info("context length %d from the checkpoint", found_value)
    return found_value

OUTTYPES = ("auto", "f32", "f16", "bf16", "nvfp4")


def outtype(value: str) -> str:
    return value.lower()


def quote_modelfile_value(text: str) -> str:
    if text == "" or any(char.isspace() for char in text) or '"' in text:
        if '"""' in text:
            raise ValueError("Modelfile value contains a triple quote, which a Modelfile cannot hold")
        if "\n" in text or '"' in text:
            return '"""' + text + '"""'
        return '"' + text + '"'
    return text


def mmproj_destinations(gguf_path: Path, sources: list[Path] | None) -> list[tuple[Path, Path]]:
    """Pair each projector with the path it will have in the output directory."""
    dest_dir = gguf_path.parent
    planned: list[tuple[Path, Path]] = []
    seen: set[str] = set()
    for source in sources or []:
        source = Path(source)
        if not source.is_file():
            raise SystemExit(f"mmproj file does not exist: {source}")
        if source.name == gguf_path.name or source.name in seen:
            raise SystemExit(
                f"mmproj {source} would be copied as {source.name}, which collides with another file in {dest_dir}."
            )
        seen.add(source.name)
        planned.append((source, dest_dir / source.name))
    return planned


def place_mmprojs(gguf_path: Path, sources: list[Path] | None, overwrite: bool) -> list[Path]:
    """Copy projector GGUFs next to the language model, keeping their file names."""
    copied: list[Path] = []
    for source, dest in mmproj_destinations(gguf_path, sources):
        try:
            same = dest.exists() and dest.resolve() == source.resolve()
        except OSError:
            same = False
        if same:
            copied.append(dest)
            continue
        if dest.exists() and not overwrite:
            raise SystemExit(f"{dest} already exists. Pass --overwrite to replace it.")
        shutil.copy2(source, dest)
        log.info("Copied vision projector to %s", dest)
        copied.append(dest)
    return copied


def refuse_existing_bundle(gguf_path: Path, overwrite: bool, mmprojs: list[Path] | None, *, skip_gguf: bool = False) -> None:
    """Stop before a long run when the GGUF, Modelfile, or a copied projector is already there."""
    if not skip_gguf and gguf_path.exists() and not overwrite:
        raise SystemExit(f"{gguf_path} already exists. Pass --overwrite to replace it.")
    modelfile = gguf_path.parent / "Modelfile"
    if modelfile.exists() and not overwrite:
        raise SystemExit(f"{modelfile} already exists. Pass --overwrite to replace it.")
    for source, dest in mmproj_destinations(gguf_path, mmprojs):
        try:
            same = dest.exists() and dest.resolve() == source.resolve()
        except OSError:
            same = False
        if dest.exists() and not same and not overwrite:
            raise SystemExit(f"{dest} already exists. Pass --overwrite to replace it.")


def _stored_value(kv: dict, name: str):
    item = kv.get(name)
    if item is None:
        return None
    return getattr(item, "value", item)


def _nonempty(value) -> bool:
    if value is None:
        return False
    if isinstance(value, str):
        return bool(value.strip())
    if isinstance(value, (list, tuple)):
        return len(value) > 0
    return True


def warn_missing_gguf_metadata(writer) -> None:
    """Warn when the GGUF metadata cannot describe a usable chat model.

    Context length, the chat template, and the tokenizer are the fields that
    still produce a file but change how it runs. Block count and embedding
    length are what llama.cpp needs to build the graph.
    """
    kv = writer.kv_data[0] if writer.kv_data else {}
    arch = writer.arch or "unknown"
    missing: list[str] = []

    context_keys = [name for name in kv if is_model_context_key(name)]
    if not context_keys:
        missing.append(f"context length ({arch}.context_length); pass --context-length to set it")
    elif all(scalar_context_length(_stored_value(kv, name)) is None for name in context_keys):
        missing.append("a single context length; pass --context-length to replace the stored value")

    template_names = [
        name for name in kv
        if name == "tokenizer.chat_template" or name == "tokenizer.chat_templates" or name.startswith("tokenizer.chat_template.")
    ]
    if not any(_nonempty(_stored_value(kv, name)) for name in template_names):
        missing.append("chat template (tokenizer.chat_template)")

    if not _nonempty(_stored_value(kv, "tokenizer.ggml.model")):
        missing.append("tokenizer model (tokenizer.ggml.model)")
    if not _nonempty(_stored_value(kv, "tokenizer.ggml.tokens")):
        missing.append("tokenizer vocabulary (tokenizer.ggml.tokens)")
    if "tokenizer.ggml.eos_token_id" not in kv:
        missing.append("end-of-sequence token (tokenizer.ggml.eos_token_id)")
    if scalar_context_length(_stored_value(kv, f"{arch}.block_count")) is None:
        missing.append(f"block count ({arch}.block_count)")
    if scalar_context_length(_stored_value(kv, f"{arch}.embedding_length")) is None:
        missing.append(f"embedding length ({arch}.embedding_length)")

    if missing:
        log.warning("GGUF is missing %s.", "; ".join(missing))


# llama.cpp enables its DeepSeek-R1 tool parser only when the template source
# contains this exact expression. The sglang DeepSeek-R1 template renders the
# same calls but does not contain the expression, so the calls stay in the
# visible reply.
_DEEPSEEK_R1_TOOL_MARKER = (
    "{{'<｜Assistant｜><｜tool▁calls▁begin｜><｜tool▁call▁begin｜>' + tool['type'] + '<｜tool▁sep｜>'"
)
# A brace was left between the two branches of that template. Jinja prints it
# after every tool call that has no text.
_DEEPSEEK_R1_STRAY_BRACE = "'<｜tool▁call▁end｜>'}}\n                }\n                {%- else %}"
_DEEPSEEK_R1_STRAY_FIXED = "'<｜tool▁call▁end｜>'}}\n                {%- else %}"


def normalize_chat_template(value):
    """Make a DeepSeek-R1 tool template recognizable to llama.cpp's parser."""
    if isinstance(value, str):
        return _normalize_one_chat_template(value)
    if isinstance(value, list):
        changed = []
        for item in value:
            if isinstance(item, str):
                changed.append(_normalize_one_chat_template(item))
            elif isinstance(item, dict) and isinstance(item.get("template"), str):
                item = dict(item)
                item["template"] = _normalize_one_chat_template(item["template"])
                changed.append(item)
            else:
                changed.append(item)
        return changed
    return value


def _normalize_one_chat_template(template: str) -> str:
    if "tool▁calls▁begin" not in template:
        return template
    if _DEEPSEEK_R1_STRAY_BRACE in template:
        template = template.replace(_DEEPSEEK_R1_STRAY_BRACE, _DEEPSEEK_R1_STRAY_FIXED, 1)
    if _DEEPSEEK_R1_TOOL_MARKER not in template:
        template = "{# " + _DEEPSEEK_R1_TOOL_MARKER + " #}\n" + template
        log.info("Adjusted the DeepSeek-R1 tool template so tool calls are parsed instead of printed.")
    return template


def flatten_chat_template(value) -> str:
    """Join a string or a list of named chat templates into one text blob."""
    if value is None:
        return ""
    if isinstance(value, bytes):
        return value.decode("utf-8", "replace")
    if isinstance(value, str):
        return value
    if isinstance(value, list):
        parts: list[str] = []
        for item in value:
            if isinstance(item, str):
                parts.append(item)
            elif isinstance(item, dict) and isinstance(item.get("template"), str):
                parts.append(item["template"])
        return "\n".join(parts)
    return ""


# Ollama accepts a tools request only when the chat template both mentions tools
# and can render a tool result. A bare "tools" string does not.
_TOOL_RESULT_MARKERS = (
    "tool_response",
    "tool_results",
    "role'] == 'tool'",
    "role'] == \"tool\"",
    "role\"] == 'tool'",
    "role\"] == \"tool\"",
    "message.role == 'tool'",
    "message.role == \"tool\"",
    "ipython",
)


def chat_template_has_tools(template: str | None) -> bool:
    """True when the template can render a tool call and the tool's result."""
    if not template:
        return False
    if "tools" not in template and "tool_call" not in template:
        return False
    if "tool_calls" not in template and "assistant_tool_call" not in template:
        return False
    return any(marker in template for marker in _TOOL_RESULT_MARKERS)


def tool_modelfile_lines(template: str | None) -> list[str]:
    """Declare tools when the stored chat template can round-trip them.

    A template that cannot render a tool result leaves the capability off.
    Ollama then rejects a request that includes tools instead of dropping the
    call on the floor.
    """
    if chat_template_has_tools(template):
        return [
            "CAPABILITY tools",
            "# Tool calls use the chat template stored in the GGUF.",
        ]
    return ["# This chat template has no tool round-trip, so tools stay off."]


def write_simple_modelfile(
    gguf_path: Path,
    overwrite: bool,
    mmprojs: list[Path] | None,
    chat_template: str | None = None,
) -> Path:
    """Write a Modelfile that points at a finished GGUF and any projector copies.

    Context is not set here. A PARAMETER num_ctx would override OLLAMA_CONTEXT_LENGTH.
    """
    modelfile = gguf_path.parent / "Modelfile"
    if modelfile.exists() and not overwrite:
        raise SystemExit(f"{modelfile} already exists. Pass --overwrite to replace it.")
    projectors = place_mmprojs(gguf_path, mmprojs, overwrite)
    lines = [f"FROM ./{gguf_path.name}"]
    for projector in projectors:
        lines.append(f"FROM ./{projector.name}")
    lines.append("")
    lines.append("# The chat template is the one stored in the GGUF.")
    lines.extend(tool_modelfile_lines(chat_template))
    if chat_template_has_tools(chat_template):
        log.info("Chat template supports tools.")
    elif chat_template:
        log.info("Chat template has no tool round-trip. Tool requests stay rejected.")
    lines.append("# Context length comes from OLLAMA_CONTEXT_LENGTH.")
    lines.append("")
    modelfile.write_text("\n".join(lines), encoding="utf-8")
    log.info("Wrote %s", modelfile)
    return modelfile