#!/usr/bin/env python3
from __future__ import annotations

import argparse
import math
import os
import re
import shutil
import struct
import sys
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path


DEFAULT_DENSE_CTX = "32768"
DEFAULT_DENSE_GPU_LAYERS = "28"
DEFAULT_MOE_CTX = "43008"
DEFAULT_MOE_GPU_LAYERS = "999"
DEFAULT_BATCH_SIZE = "2048"
DEFAULT_MOE_BATCH_SIZE = "1024"
DEFAULT_UBATCH_SIZE = "512"
DEFAULT_PARALLEL = "4"
DEFAULT_MOE_PARALLEL = "4"
DEFAULT_ASR_PARALLEL = "1"
DEFAULT_MIN_CTX = 512
DEFAULT_CACHE_TYPE = "f16"
LONG_CONTEXT_CACHE_TYPE = "q8_0"
QUANT_PRIORITY = {
    "Q4_K_M": 0,
    "Q5_K_M": 1,
    "Q6_K": 2,
    "Q8_0": 3,
    "Q4_K_S": 4,
    "Q5_K_S": 5,
    "Q4_0": 6,
    "Q4_1": 7,
    "Q5_0": 8,
    "Q5_1": 9,
    "Q3_K_M": 10,
    "Q3_K_L": 11,
    "Q2_K": 12,
    "F16": 90,
    "BF16": 91,
}
GGUF_MAGIC = b"GGUF"
GGUF_TYPE_UINT8 = 0
GGUF_TYPE_INT8 = 1
GGUF_TYPE_UINT16 = 2
GGUF_TYPE_INT16 = 3
GGUF_TYPE_UINT32 = 4
GGUF_TYPE_INT32 = 5
GGUF_TYPE_FLOAT32 = 6
GGUF_TYPE_BOOL = 7
GGUF_TYPE_STRING = 8
GGUF_TYPE_ARRAY = 9
GGUF_TYPE_UINT64 = 10
GGUF_TYPE_INT64 = 11
GGUF_TYPE_FLOAT64 = 12
GGUF_SCALAR_FORMATS = {
    GGUF_TYPE_UINT8: "<B",
    GGUF_TYPE_INT8: "<b",
    GGUF_TYPE_UINT16: "<H",
    GGUF_TYPE_INT16: "<h",
    GGUF_TYPE_UINT32: "<I",
    GGUF_TYPE_INT32: "<i",
    GGUF_TYPE_FLOAT32: "<f",
    GGUF_TYPE_BOOL: "<?",
    GGUF_TYPE_UINT64: "<Q",
    GGUF_TYPE_INT64: "<q",
    GGUF_TYPE_FLOAT64: "<d",
}
MOE_CPU_OFFLOAD_LARGE_THRESHOLD_GIB = 20.0
@dataclass(frozen=True)
class ModelMetadata:
    architecture: str | None
    context_length: int | None
    block_count: int | None
    expert_count: int | None
    expert_used_count: int | None
    size_label: str | None
    general_name: str | None
    file_size_gib: float


@dataclass(frozen=True)
class ModelEntry:
    publisher: str
    model_dir: str
    model_path: Path
    mmproj_path: Path | None
    quant_key: str
    metadata: ModelMetadata


@dataclass(frozen=True)
class GenerateResult:
    output_path: Path
    backup_path: Path | None
    entry_count: int


@dataclass(frozen=True)
class ValueRange:
    minimum: int
    maximum: int


def is_mmproj_file(path: Path) -> bool:
    name = path.name.lower()
    stem = path.stem.lower()
    return (
        name.startswith("mmproj")
        or ".mmproj-" in name
        or stem.endswith(".mmproj")
        or stem.startswith("mmproj")
    )


def normalize_slug(value: str) -> str:
    lowered = value.lower()
    normalized = re.sub(r"[^a-z0-9]+", "-", lowered)
    normalized = re.sub(r"-{2,}", "-", normalized)
    return normalized.strip("-") or "model"


def detect_quant_key(filename: str) -> str:
    upper_name = filename.upper()
    for quant_key in QUANT_PRIORITY:
        if quant_key in upper_name:
            return quant_key
    matches = re.findall(r"(Q\d(?:_[A-Z0-9]+)+|Q\d+|F16|BF16)", upper_name)
    return matches[0] if matches else "UNKNOWN"


def normalize_quant_key(quant_key: str) -> str:
    return normalize_slug(quant_key)


def quant_rank(quant_key: str) -> tuple[int, str]:
    return (QUANT_PRIORITY.get(quant_key, 50), quant_key)


def read_gguf_string(handle) -> str:
    length = struct.unpack("<Q", handle.read(8))[0]
    return handle.read(length).decode("utf-8", errors="replace")


def read_gguf_scalar(handle, value_type: int):
    fmt = GGUF_SCALAR_FORMATS.get(value_type)
    if fmt is None:
        raise ValueError(f"Unsupported GGUF scalar type: {value_type}")
    size = struct.calcsize(fmt)
    return struct.unpack(fmt, handle.read(size))[0]


def read_gguf_value(handle, value_type: int):
    if value_type == GGUF_TYPE_STRING:
        return read_gguf_string(handle)
    if value_type == GGUF_TYPE_ARRAY:
        item_type = struct.unpack("<I", handle.read(4))[0]
        item_count = struct.unpack("<Q", handle.read(8))[0]
        return [read_gguf_value(handle, item_type) for _ in range(item_count)]
    return read_gguf_scalar(handle, value_type)


def parse_gguf_metadata(model_path: Path) -> dict[str, object]:
    metadata: dict[str, object] = {}
    with model_path.open("rb") as handle:
        magic = handle.read(4)
        if magic != GGUF_MAGIC:
            raise ValueError(f"Not a GGUF file: {model_path}")

        version = struct.unpack("<I", handle.read(4))[0]
        if version not in {2, 3}:
            raise ValueError(f"Unsupported GGUF version {version}: {model_path}")

        tensor_count = struct.unpack("<Q", handle.read(8))[0]
        metadata_count = struct.unpack("<Q", handle.read(8))[0]

        for _ in range(metadata_count):
            key = read_gguf_string(handle)
            value_type = struct.unpack("<I", handle.read(4))[0]
            metadata[key] = read_gguf_value(handle, value_type)

        # Stop after metadata. Tensor table is not needed.
        _ = tensor_count

    return metadata


def parse_param_billions(value: str) -> float | None:
    match = re.search(r"(\d+(?:\.\d+)?)\s*([kmbtq])", value.strip().lower())
    if not match:
        return None
    amount = float(match.group(1))
    scale = match.group(2)
    scale_map = {
        "k": 1e-6,
        "m": 1e-3,
        "b": 1.0,
        "t": 1e3,
        "q": 1e6,
    }
    return amount * scale_map[scale]


def estimate_param_billions_from_name(value: str) -> float | None:
    match = re.search(r"(\d+(?:\.\d+)?)\s*b", value.lower())
    if not match:
        return None
    return float(match.group(1))


def clamp_int(value: int, minimum: int, maximum: int) -> int:
    return max(minimum, min(value, maximum))


def get_logical_cpu_count() -> int:
    return os.cpu_count() or 8


def build_core_fields(*, batch_size: str, ubatch_size: str, parallel: str, mmap: str) -> dict[str, str]:
    return {
        "batch-size": batch_size,
        "ubatch-size": ubatch_size,
        "parallel": parallel,
        "kv-unified": "true",
        "kv-offload": "true",
        "mlock": "true",
        "mmap": mmap,
        "seed": "-1",
        "flash-attn": "true",
    }


def build_sampling_defaults(family: str) -> dict[str, str]:
    if family == "moe":
        return {
            "temp": "0.7",
            "top-k": "40",
            "top-p": "0.95",
            "min-p": "0.0",
            "repeat-penalty": "1.0",
        }
    return {
        "temp": "0.8",
        "top-k": "40",
        "top-p": "0.95",
        "min-p": "0.05",
        "repeat-penalty": "1.0",
    }


def build_model_metadata(model_path: Path) -> ModelMetadata:
    file_size_gib = model_path.stat().st_size / (1024 ** 3)
    try:
        raw = parse_gguf_metadata(model_path)
    except Exception:
        raw = {}

    architecture = raw.get("general.architecture")
    if not isinstance(architecture, str):
        architecture = None

    def get_arch_value(suffix: str) -> object | None:
        if architecture is None:
            return None
        return raw.get(f"{architecture}.{suffix}")

    context_length = get_arch_value("context_length")
    block_count = get_arch_value("block_count")
    expert_count = get_arch_value("expert_count")
    expert_used_count = get_arch_value("expert_used_count")
    size_label = raw.get("general.size_label")
    general_name = raw.get("general.name")

    return ModelMetadata(
        architecture=architecture,
        context_length=context_length if isinstance(context_length, int) else None,
        block_count=block_count if isinstance(block_count, int) else None,
        expert_count=expert_count if isinstance(expert_count, int) else None,
        expert_used_count=expert_used_count if isinstance(expert_used_count, int) else None,
        size_label=size_label if isinstance(size_label, str) else None,
        general_name=general_name if isinstance(general_name, str) else None,
        file_size_gib=file_size_gib,
    )


def classify_model_family(entry: ModelEntry) -> str:
    architecture = (entry.metadata.architecture or "").lower()
    if entry.metadata.expert_count and entry.metadata.expert_count > 0:
        return "moe"
    if "moe" in architecture:
        return "moe"
    if entry.mmproj_path is not None:
        return "asr_multimodal"
    return "dense"


def classify_moe_tier(entry: ModelEntry) -> str:
    metadata = entry.metadata
    score = 0
    if metadata.expert_count:
        score += 2
    if metadata.expert_used_count and metadata.expert_used_count >= 4:
        score += 1
    if metadata.file_size_gib >= MOE_CPU_OFFLOAD_LARGE_THRESHOLD_GIB:
        score += 2
    elif metadata.file_size_gib >= 12.0:
        score += 1
    if metadata.context_length and metadata.context_length >= 131072:
        score += 2
    elif metadata.context_length and metadata.context_length >= 65536:
        score += 1
    if metadata.block_count and metadata.block_count >= 80:
        score += 1

    if score >= 5:
        return "large"
    if score >= 3:
        return "medium"
    return "small"


def choose_gpu_layers(entry: ModelEntry, fallback: str) -> str:
    block_count = entry.metadata.block_count or 0
    if block_count > 0:
        return str(block_count)
    return fallback


def choose_context_size(entry: ModelEntry, family: str) -> str:
    max_context = entry.metadata.context_length or 0
    if max_context <= 0:
        defaults = {
            "dense": DEFAULT_DENSE_CTX,
            "asr_multimodal": "8192",
            "moe": DEFAULT_MOE_CTX,
        }
        return defaults[family]

    target_map = {
        "dense": int(DEFAULT_DENSE_CTX),
        "asr_multimodal": 8192,
        "moe": int(DEFAULT_MOE_CTX),
    }
    target = min(max_context, target_map[family])
    chosen = clamp_int(target, 1 if max_context < DEFAULT_MIN_CTX else DEFAULT_MIN_CTX, max_context)
    return str(chosen)


def choose_cache_type(entry: ModelEntry, family: str) -> str:
    context_length = entry.metadata.context_length or 0
    if family == "moe" or context_length >= 65536:
        return LONG_CONTEXT_CACHE_TYPE
    return DEFAULT_CACHE_TYPE


def choose_moe_cpu_layers(entry: ModelEntry) -> str | None:
    block_count = entry.metadata.block_count or 0
    if block_count <= 0:
        return None
    return "0"


def build_dense_values(profile: str) -> dict[str, str]:
    _ = profile
    values = build_core_fields(
        batch_size=DEFAULT_BATCH_SIZE,
        ubatch_size=DEFAULT_UBATCH_SIZE,
        parallel=DEFAULT_PARALLEL,
        mmap="true",
    )
    values["n-gpu-layers"] = DEFAULT_DENSE_GPU_LAYERS
    values["threads"] = str(default_thread_pool_size())
    values["cache-type-k"] = DEFAULT_CACHE_TYPE
    values["cache-type-v"] = DEFAULT_CACHE_TYPE
    values.update(build_sampling_defaults("dense"))
    return values


def build_asr_multimodal_values(entry: ModelEntry, profile: str) -> dict[str, str]:
    _ = profile
    values = build_core_fields(
        batch_size=DEFAULT_MOE_BATCH_SIZE,
        ubatch_size=DEFAULT_UBATCH_SIZE,
        parallel=DEFAULT_ASR_PARALLEL,
        mmap="true",
    )
    values["n-gpu-layers"] = choose_gpu_layers(entry, DEFAULT_DENSE_GPU_LAYERS)
    values["threads"] = str(default_thread_pool_size())
    values["ctx-size"] = choose_context_size(entry, "asr_multimodal")
    cache_type = choose_cache_type(entry, "asr_multimodal")
    values["cache-type-k"] = cache_type
    values["cache-type-v"] = cache_type
    values.update(build_sampling_defaults("dense"))
    return values


def build_moe_values(entry: ModelEntry, profile: str) -> dict[str, str]:
    _ = profile
    values = build_core_fields(
        batch_size=DEFAULT_BATCH_SIZE,
        ubatch_size=DEFAULT_UBATCH_SIZE,
        parallel=DEFAULT_MOE_PARALLEL,
        mmap="true",
    )
    values["ctx-size"] = choose_context_size(entry, "moe")
    values["n-gpu-layers"] = choose_gpu_layers(entry, DEFAULT_MOE_GPU_LAYERS)
    values["threads"] = str(default_thread_pool_size())
    values["cache-type-k"] = LONG_CONTEXT_CACHE_TYPE
    values["cache-type-v"] = LONG_CONTEXT_CACHE_TYPE
    values.update(build_sampling_defaults("moe"))
    n_cpu_moe = choose_moe_cpu_layers(entry)
    if n_cpu_moe is not None:
        values["n-cpu-moe"] = n_cpu_moe
    return values


def default_thread_pool_size() -> int:
    logical_cpu_count = get_logical_cpu_count()
    return max(1, math.floor(logical_cpu_count * 0.7))


def scan_models(models_dir: Path, select_mode: str) -> list[ModelEntry]:
    models_dir = Path(models_dir)
    if not models_dir.is_dir():
        raise FileNotFoundError(f"Models directory not found: {models_dir}")

    grouped: dict[tuple[str, str], list[ModelEntry]] = {}
    for path in sorted(models_dir.rglob("*.gguf")):
        if is_mmproj_file(path):
            continue

        relative_parts = path.relative_to(models_dir).parts
        if len(relative_parts) < 3:
            publisher = relative_parts[0] if relative_parts else "unknown"
            model_dir = path.parent.name
        else:
            publisher = relative_parts[0]
            model_dir = relative_parts[1]

        mmproj_candidates = sorted(candidate for candidate in path.parent.glob("*.gguf") if is_mmproj_file(candidate))
        mmproj_path = mmproj_candidates[0] if mmproj_candidates else None

        entry = ModelEntry(
            publisher=publisher,
            model_dir=model_dir,
            model_path=path,
            mmproj_path=mmproj_path,
            quant_key=detect_quant_key(path.name),
            metadata=build_model_metadata(path),
        )
        grouped.setdefault((publisher, model_dir), []).append(entry)

    selected: list[ModelEntry] = []
    for _, entries in sorted(grouped.items()):
        if select_mode == "all":
            selected.extend(sorted(entries, key=lambda item: item.model_path.name.lower()))
        elif select_mode == "best":
            best_entry = min(
                entries,
                key=lambda item: (quant_rank(item.quant_key), item.model_path.name.lower()),
            )
            selected.append(best_entry)
        else:
            raise ValueError(f"Unsupported select mode: {select_mode}")

    if not selected:
        raise ValueError(f"No model GGUF files found under {models_dir}")
    return selected


def build_preset_values(entry: ModelEntry, profile: str) -> dict[str, str]:
    if profile not in {"conservative", "smart"}:
        raise ValueError(f"Unsupported profile: {profile}")

    family = classify_model_family(entry)
    if family == "asr_multimodal":
        return build_asr_multimodal_values(entry, profile)
    if family == "moe":
        return build_moe_values(entry, profile)
    values = build_dense_values(profile)
    values["ctx-size"] = choose_context_size(entry, "dense")
    values["n-gpu-layers"] = choose_gpu_layers(entry, values["n-gpu-layers"])
    cache_type = choose_cache_type(entry, "dense")
    values["cache-type-k"] = cache_type
    values["cache-type-v"] = cache_type
    return values


def build_value_ranges(entry: ModelEntry, values: dict[str, str]) -> dict[str, ValueRange]:
    ranges: dict[str, ValueRange] = {}
    context_length = entry.metadata.context_length or 0
    block_count = entry.metadata.block_count or 0
    thread_count = default_thread_pool_size()

    if "ctx-size" in values and context_length > 0:
        ranges["ctx-size"] = ValueRange(minimum=min(DEFAULT_MIN_CTX, context_length), maximum=context_length)
    if "n-gpu-layers" in values and block_count > 0:
        ranges["n-gpu-layers"] = ValueRange(minimum=0, maximum=block_count)
    if "threads" in values:
        ranges["threads"] = ValueRange(minimum=1, maximum=get_logical_cpu_count() or thread_count)
    if "batch-size" in values:
        ranges["batch-size"] = ValueRange(minimum=32, maximum=int(values["batch-size"]))
    if "ubatch-size" in values:
        ranges["ubatch-size"] = ValueRange(minimum=32, maximum=int(values["batch-size"]))
    if "parallel" in values:
        ranges["parallel"] = ValueRange(minimum=1, maximum=int(values["parallel"]))
    if "n-cpu-moe" in values and block_count > 0:
        ranges["n-cpu-moe"] = ValueRange(minimum=0, maximum=block_count)

    return ranges


def format_value_comment(key: str, value: str, ranges: dict[str, ValueRange]) -> str | None:
    value_range = ranges.get(key)
    if value_range is None:
        return None
    return f"; {key} range = {value_range.minimum}..{value_range.maximum}; chosen = {value}"


def make_alias(
    publisher: str,
    model_dir: str,
    model_file: str,
    alias_style: str,
    used_aliases: set[str],
) -> str:
    if alias_style == "filename":
        base = Path(model_file).stem
    elif alias_style == "section":
        base = f"{publisher}/{model_dir}:{detect_quant_key(model_file)}"
    elif alias_style == "short":
        base = Path(model_file).stem
    else:
        raise ValueError(f"Unsupported alias style: {alias_style}")

    if alias_style == "section":
        alias = base
    else:
        alias = normalize_slug(base)
    if alias_style == "short":
        alias = re.sub(r"-gguf$", "", alias)

    unique_alias = alias
    suffix = 2
    while unique_alias in used_aliases:
        if alias_style == "section":
            unique_alias = f"{alias}-{suffix}"
        else:
            unique_alias = f"{alias}-{suffix}"
        suffix += 1
    used_aliases.add(unique_alias)
    return unique_alias


def render_ini(
    entries: list[ModelEntry],
    models_dir: Path,
    select_mode: str,
    profile: str,
    alias_style: str,
    backup_path: Path | None,
) -> str:
    generated_at = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    lines = [
        "; Auto-generated llama.cpp models preset",
        f"; generated-at = {generated_at}",
        f"; models-dir = {models_dir}",
        f"; select = {select_mode}",
        f"; profile = {profile}",
        f"; alias-style = {alias_style}",
        f"; backup = {backup_path if backup_path else 'none'}",
        "",
        "version = 1",
        "",
    ]

    used_aliases: set[str] = set()
    for entry in entries:
        alias = make_alias(
            publisher=entry.publisher,
            model_dir=entry.model_dir,
            model_file=entry.model_path.name,
            alias_style=alias_style,
            used_aliases=used_aliases,
        )
        values = build_preset_values(entry=entry, profile=profile)

        lines.append(f"[{alias}]")
        lines.append(f"model = {entry.model_path}")
        if entry.mmproj_path is not None:
            lines.append(f"mmproj = {entry.mmproj_path}")
        lines.append(f"alias = {alias}")
        ranges = build_value_ranges(entry, values)
        ordered_keys = [
            "ctx-size",
            "n-gpu-layers",
            "threads",
            "batch-size",
            "ubatch-size",
            "parallel",
            "temp",
            "top-k",
            "top-p",
            "min-p",
            "repeat-penalty",
            "cache-type-k",
            "cache-type-v",
            "kv-unified",
            "kv-offload",
            "mlock",
            "mmap",
            "seed",
            "flash-attn",
        ]
        optional_keys = ["n-cpu-moe", "cpu-moe"]
        sampling_keys = {"temp", "top-k", "top-p", "min-p", "repeat-penalty"}
        sampling_comment_written = False
        for key in ordered_keys:
            if key in sampling_keys and not sampling_comment_written:
                lines.append("; Heuristic sampling defaults / 经验采样默认值, not GGUF metadata-derived")
                sampling_comment_written = True
            comment = format_value_comment(key, values[key], ranges)
            if comment is not None:
                lines.append(comment)
            lines.append(f"{key} = {values[key]}")
        for key in optional_keys:
            if key in values:
                comment = format_value_comment(key, values[key], ranges)
                if comment is not None:
                    lines.append(comment)
                lines.append(f"{key} = {values[key]}")
        lines.append("")

    return "\n".join(lines).rstrip() + "\n"


def maybe_backup_file(output_path: Path, backup: bool) -> Path | None:
    if not backup or not output_path.exists():
        return None
    timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
    backup_path = output_path.with_name(f"{output_path.stem}.{timestamp}.bak{output_path.suffix}")
    shutil.copy2(output_path, backup_path)
    return backup_path


def generate_preset_file(
    models_dir: Path,
    output_path: Path,
    select_mode: str,
    profile: str,
    alias_style: str,
    backup: bool,
) -> GenerateResult:
    entries = scan_models(models_dir=models_dir, select_mode=select_mode)
    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    backup_path = maybe_backup_file(output_path=output_path, backup=backup)
    ini_text = render_ini(
        entries=entries,
        models_dir=Path(models_dir),
        select_mode=select_mode,
        profile=profile,
        alias_style=alias_style,
        backup_path=backup_path,
    )
    output_path.write_text(ini_text, encoding="utf-8")
    return GenerateResult(
        output_path=output_path,
        backup_path=backup_path,
        entry_count=len(entries),
    )


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Scan an LM Studio-style models directory and generate models-preset.ini"
    )
    parser.add_argument("--models-dir", required=True, help="LM Studio-style models root directory")
    parser.add_argument("--output", default="models-preset.ini", help="Output INI file path")
    parser.add_argument("--select", choices=["all", "best"], default="all")
    parser.add_argument("--profile", choices=["conservative", "smart"], default="smart")
    parser.add_argument("--alias-style", choices=["filename", "section", "short"], default="section")
    parser.add_argument("--backup", dest="backup", action="store_true", default=False)
    parser.add_argument("--no-backup", dest="backup", action="store_false")
    return parser


def main(argv: list[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)

    try:
        result = generate_preset_file(
            models_dir=Path(args.models_dir),
            output_path=Path(args.output),
            select_mode=args.select,
            profile=args.profile,
            alias_style=args.alias_style,
            backup=args.backup,
        )
    except Exception as exc:
        print(f"ERROR: {exc}", file=sys.stderr)
        return 1

    print(f"Generated {result.entry_count} model preset section(s) at: {result.output_path}")
    if result.backup_path is not None:
        print(f"Backup created: {result.backup_path}")
    return 0


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