#!/usr/bin/env python3
import argparse
import json
import re
import struct
from datetime import datetime
from pathlib import Path


DEFAULT_MODELS_ROOT = Path.home() / ".lmstudio" / "models"
DEFAULT_OUTPUT_ROOT = Path.home() / ".lmstudio" / "hub" / "models"
DEFAULT_INTERNAL_ROOT = Path.home() / ".lmstudio" / ".internal"
DEFAULT_GGUF_METADATA_CACHE = DEFAULT_INTERNAL_ROOT / "gguf-metadata-cache.json"
DEFAULT_MODEL_INDEX_CACHE = DEFAULT_INTERNAL_ROOT / "model-index-cache.json"

GGUF_VALUE_TYPES = {
    0: ("<B", 1),
    1: ("<b", 1),
    2: ("<H", 2),
    3: ("<h", 2),
    4: ("<I", 4),
    5: ("<i", 4),
    6: ("<f", 4),
    7: ("<?", 1),
    10: ("<Q", 8),
    11: ("<q", 8),
    12: ("<d", 8),
}


def _bool_text(value: bool) -> str:
    return "true" if value else "false"


def normalize_owner(owner: str) -> str:
    return _slugify(owner)


def _read_struct(data: bytes, offset: int, fmt: str) -> tuple[object, int]:
    size = struct.calcsize(fmt)
    if offset + size > len(data):
        raise ValueError("Unexpected end of GGUF metadata")
    value = struct.unpack_from(fmt, data, offset)[0]
    return value, offset + size


def _read_gguf_string(data: bytes, offset: int) -> tuple[str, int]:
    length, offset = _read_struct(data, offset, "<Q")
    end = offset + int(length)
    if end > len(data):
        raise ValueError("Invalid GGUF string length")
    return data[offset:end].decode("utf-8", errors="replace"), end


def _read_gguf_value(data: bytes, offset: int, value_type: int) -> tuple[object, int]:
    if value_type == 8:
        return _read_gguf_string(data, offset)
    if value_type == 9:
        nested_type, offset = _read_struct(data, offset, "<I")
        count, offset = _read_struct(data, offset, "<Q")
        values = []
        for _ in range(int(count)):
            item, offset = _read_gguf_value(data, offset, int(nested_type))
            values.append(item)
        return values, offset
    if value_type not in GGUF_VALUE_TYPES:
        raise ValueError(f"Unsupported GGUF metadata value type: {value_type}")
    fmt, _ = GGUF_VALUE_TYPES[value_type]
    return _read_struct(data, offset, fmt)


def _normalize_cached_metadata(metadata: dict[str, object]) -> dict[str, object]:
    normalized: dict[str, object] = {}
    architecture = metadata.get("arch")
    if architecture is not None:
        normalized["general.architecture"] = architecture
    name = metadata.get("name")
    if name is not None:
        normalized["general.name"] = name
    chat_template = metadata.get("chatTemplate")
    if chat_template is not None:
        normalized["tokenizer.chat_template"] = chat_template
    parameters = metadata.get("parameters")
    if parameters is not None:
        normalized["general.parameter_count_label"] = str(parameters)
    context_length = metadata.get("contextLength")
    if context_length is not None:
        normalized["general.context_length"] = str(context_length)
    bos_token = metadata.get("bosToken")
    if bos_token is not None:
        normalized["tokenizer.ggml.bos_token"] = bos_token
    eos_token = metadata.get("eosToken")
    if eos_token is not None:
        normalized["tokenizer.ggml.eos_token"] = eos_token
    return normalized


def load_gguf_metadata_cache(cache_path: Path = DEFAULT_GGUF_METADATA_CACHE) -> dict[str, dict[str, object]]:
    if not cache_path.exists():
        return {}
    payload = json.loads(cache_path.read_text(encoding="utf-8"))
    entries = payload.get("json", {}).get("map", [])
    cache_index: dict[str, dict[str, object]] = {}
    for entry in entries:
        if not isinstance(entry, list) or len(entry) != 2:
            continue
        file_path, cached_payload = entry
        if not isinstance(file_path, str) or not isinstance(cached_payload, dict):
            continue
        metadata = cached_payload.get("metadata")
        if isinstance(metadata, dict):
            cache_index[file_path] = _normalize_cached_metadata(metadata)
    return cache_index


def _read_exact(file_obj, size: int) -> bytes:
    data = file_obj.read(size)
    if len(data) != size:
        raise ValueError("Unexpected end of GGUF metadata stream")
    return data


def _read_struct_stream(file_obj, fmt: str) -> object:
    size = struct.calcsize(fmt)
    return struct.unpack(fmt, _read_exact(file_obj, size))[0]


def _read_gguf_string_stream(file_obj) -> str:
    length = int(_read_struct_stream(file_obj, "<Q"))
    return _read_exact(file_obj, length).decode("utf-8", errors="replace")


def _read_gguf_value_stream(file_obj, value_type: int) -> object:
    if value_type == 8:
        return _read_gguf_string_stream(file_obj)
    if value_type == 9:
        nested_type = int(_read_struct_stream(file_obj, "<I"))
        count = int(_read_struct_stream(file_obj, "<Q"))
        return [_read_gguf_value_stream(file_obj, nested_type) for _ in range(count)]
    if value_type not in GGUF_VALUE_TYPES:
        raise ValueError(f"Unsupported GGUF metadata value type: {value_type}")
    fmt, _ = GGUF_VALUE_TYPES[value_type]
    return _read_struct_stream(file_obj, fmt)


def read_gguf_metadata_streaming(path: Path) -> dict[str, object]:
    with path.open("rb") as file_obj:
        if _read_exact(file_obj, 4) != b"GGUF":
            raise ValueError(f"Not a GGUF file: {path}")

        version = int(_read_struct_stream(file_obj, "<I"))
        if version not in (2, 3):
            raise ValueError(f"Unsupported GGUF version: {version}")

        _ = _read_struct_stream(file_obj, "<Q")
        metadata_count = int(_read_struct_stream(file_obj, "<Q"))

        metadata: dict[str, object] = {}
        for _ in range(metadata_count):
            key = _read_gguf_string_stream(file_obj)
            value_type = int(_read_struct_stream(file_obj, "<I"))
            metadata[key] = _read_gguf_value_stream(file_obj, value_type)
        return metadata


def resolve_gguf_metadata(
    gguf_path: Path,
    cache_index: dict[str, dict[str, object]],
) -> dict[str, object]:
    resolved = gguf_path.resolve().as_posix()
    cached = cache_index.get(resolved)
    if cached is not None:
        return dict(cached)
    return read_gguf_metadata_streaming(gguf_path)


def _slugify(text: str) -> str:
    lowered = text.strip().lower()
    lowered = lowered.replace("_", "-")
    lowered = re.sub(r"(?<=[a-z0-9])(?=[A-Z])", "-", lowered)
    lowered = lowered.replace("qwen3vl", "qwen3-vl")
    lowered = lowered.replace("gemma4", "gemma-4")
    lowered = re.sub(r"[^a-z0-9.]+", "-", lowered)
    lowered = re.sub(r"-{2,}", "-", lowered)
    return lowered.strip("-")


def _extract_context_length(metadata: dict[str, object]) -> int | None:
    for key in (
        "qwen2.context_length",
        "qwen.context_length",
        "llama.context_length",
        "gemma.context_length",
        "general.context_length",
    ):
        value = metadata.get(key)
        if value is None:
            continue
        text = str(value).strip()
        if text.isdigit():
            return int(text)
    return None


def _extract_params_string(metadata: dict[str, object], base_name: str) -> str | None:
    for key in ("general.size_label", "general.parameter_count_label", "general.basename"):
        value = metadata.get(key)
        if isinstance(value, str):
            match = re.search(r"(\d+(?:\.\d+)?B(?:-A\d+B)?)", value, re.IGNORECASE)
            if match:
                raw = match.group(1).upper()
                return raw.split("-A")[0] if "-A" in raw else raw
    match = re.search(r"(\d+(?:\.\d+)?B)(?:-A\d+B)?", base_name, re.IGNORECASE)
    if match:
        return match.group(1).upper()
    return None


def _extract_base_name(metadata: dict[str, object]) -> str:
    for key in ("general.basename", "general.name"):
        value = metadata.get(key)
        if isinstance(value, str) and value.strip():
            text = value.strip()
            if text.upper().endswith("-GGUF"):
                text = text[:-5]
            return text
    return "local-model"


def _detect_capabilities(chat_template: str) -> dict[str, bool]:
    lowered = chat_template.lower()
    has_reasoning = "enable_thinking" in lowered or "preserve_thinking" in lowered
    has_tools = any(
        token in lowered
        for token in (
            "<tools>",
            "<tool_call>",
            "<tool_response>",
            "<|tool>",
            "<tool|>",
            "<|tool|>",
            "<|tool_call>",
            "<tool_call|>",
            "<|tool_call|>",
            "<|tool_response>",
            "<tool_response|>",
            "<|tool_response|>",
        )
    )
    has_vision = any(
        token in lowered
        for token in (
            "<|vision_start|>",
            "<image_soft_token>",
            "<audio_soft_token>",
            "<|image_pad|>",
            "<|video_pad|>",
        )
    )
    return {
        "reasoning": has_reasoning,
        "vision": has_vision,
        "tools": has_tools,
    }


def _architecture_defaults(architecture: str) -> dict[str, object]:
    if architecture == "qwen35moe":
        return {
            "temperature": 0.6,
            "top_k": 20,
            "top_p": 0.95,
            "min_p_checked": False,
            "repeat_penalty_checked": False,
            "presence_penalty_checked": False,
            "presence_penalty": 0.0,
            "reasoning_parsing": None,
            "supports_preserve_thinking": True,
        }
    if architecture == "qwen35":
        return {
            "temperature": 1.0,
            "top_k": 20,
            "top_p": 0.95,
            "min_p_checked": False,
            "repeat_penalty_checked": False,
            "presence_penalty_checked": True,
            "presence_penalty": 1.5,
            "reasoning_parsing": None,
            "supports_preserve_thinking": False,
        }
    if architecture == "gemma4":
        return {
            "temperature": 1.0,
            "top_k": 64,
            "top_p": 0.95,
            "min_p_checked": None,
            "repeat_penalty_checked": None,
            "presence_penalty_checked": None,
            "presence_penalty": None,
            "reasoning_parsing": {
                "enabled": True,
                "startString": "<|channel>thought",
                "endString": "<channel|>",
            },
            "supports_preserve_thinking": False,
        }
    return {
        "temperature": 1.0,
        "top_k": 40,
        "top_p": 0.95,
        "min_p_checked": None,
        "repeat_penalty_checked": None,
        "presence_penalty_checked": None,
        "presence_penalty": None,
        "reasoning_parsing": None,
        "supports_preserve_thinking": False,
    }


def _has_sibling_mmproj(gguf_path: Path) -> bool:
    for sibling in gguf_path.parent.glob("*.gguf"):
        if sibling == gguf_path:
            continue
        if "mmproj" in sibling.name.lower():
            return True
    return False


def collect_directory_capabilities(gguf_path: Path) -> dict[str, bool]:
    return {
        "has_sibling_mmproj": _has_sibling_mmproj(gguf_path),
    }


def apply_directory_capability_overrides(
    capabilities: dict[str, bool],
    directory_capabilities: dict[str, bool],
) -> dict[str, bool]:
    augmented = dict(capabilities)
    if directory_capabilities.get("has_sibling_mmproj"):
        augmented["vision"] = True
    return augmented


def build_virtual_model_profile(
    metadata: dict[str, object],
    *,
    has_sibling_mmproj: bool = False,
) -> dict[str, object]:
    architecture = str(
        metadata.get("general.architecture")
        or metadata.get("general.arch")
        or metadata.get("architecture")
        or ""
    ).strip()
    chat_template = str(metadata.get("tokenizer.chat_template") or "")
    base_name = _extract_base_name(metadata)
    basename_slug = _slugify(base_name)
    metadata_capabilities = _detect_capabilities(chat_template)
    capabilities = apply_directory_capability_overrides(
        metadata_capabilities,
        {"has_sibling_mmproj": has_sibling_mmproj},
    )
    defaults = _architecture_defaults(architecture)
    reasoning_enabled = capabilities["reasoning"]
    needs_enable_thinking = reasoning_enabled and "enable_thinking" in chat_template.lower()
    needs_preserve_thinking = needs_enable_thinking and bool(defaults["supports_preserve_thinking"])

    return {
        "architecture": architecture,
        "base_name": base_name,
        "basename_slug": basename_slug,
        "params_string": _extract_params_string(metadata, base_name),
        "context_length": _extract_context_length(metadata),
        "chat_template": chat_template,
        "capabilities": {
            "reasoning": reasoning_enabled,
            "vision": capabilities["vision"],
            "tools": capabilities["tools"],
        },
        "needs_enable_thinking": needs_enable_thinking,
        "needs_preserve_thinking": needs_preserve_thinking,
        "defaults": defaults,
    }


def infer_virtual_model_slug(directory_name: str, metadata: dict[str, object]) -> str:
    base_name = directory_name.strip() or _extract_base_name(metadata)
    if base_name.upper().endswith("-GGUF"):
        base_name = base_name[:-5]
    return _slugify(base_name)


def infer_model_key(models_root: Path, gguf_path: Path) -> str:
    return "/".join(gguf_path.relative_to(models_root).parts)


def _build_base_block(model_key: str) -> list[str]:
    return [
        "base:",
        f"  - key: {model_key}",
        "    sources: []",
    ]


def _append_metadata_overrides(lines: list[str], profile: dict[str, object]) -> None:
    capabilities = profile["capabilities"]
    lines.extend(
        [
            "metadataOverrides:",
            "  domain: llm",
            "  architectures:",
            f"    - {profile['architecture']}",
            "  compatibilityTypes:",
            "    - gguf",
        ]
    )
    if profile["params_string"]:
        lines.extend(
            [
                "  paramsStrings:",
                f"    - {profile['params_string']}",
            ]
        )
    if profile["context_length"]:
        lines.extend(
            [
                "  contextLengths:",
                f"    - {profile['context_length']}",
            ]
        )
    lines.append(f"  vision: {_bool_text(bool(capabilities['vision']))}")
    lines.append(f"  reasoning: {_bool_text(bool(capabilities['reasoning']))}")
    lines.append(f"  trainedForToolUse: {_bool_text(bool(capabilities['tools']))}")


def _append_operation_fields(lines: list[str], profile: dict[str, object]) -> None:
    defaults = profile["defaults"]
    lines.extend(
        [
            "config:",
            "  operation:",
            "    fields:",
            "      - key: llm.prediction.temperature",
            f"        value: {defaults['temperature']}",
            "      - key: llm.prediction.topKSampling",
            f"        value: {defaults['top_k']}",
            "      - key: llm.prediction.topPSampling",
            "        value:",
            "          checked: true",
            f"          value: {defaults['top_p']}",
        ]
    )
    if defaults["min_p_checked"] is not None:
        lines.extend(
            [
                "      - key: llm.prediction.minPSampling",
                "        value:",
                f"          checked: {_bool_text(bool(defaults['min_p_checked']))}",
                "          value: 0",
            ]
        )
    if defaults["repeat_penalty_checked"] is not None:
        lines.extend(
            [
                "      - key: llm.prediction.repeatPenalty",
                "        value:",
                f"          checked: {_bool_text(bool(defaults['repeat_penalty_checked']))}",
                "          value: 1.0",
            ]
        )
    if defaults["presence_penalty_checked"] is not None:
        lines.extend(
            [
                "      - key: llm.prediction.llama.presencePenalty",
                "        value:",
                f"          checked: {_bool_text(bool(defaults['presence_penalty_checked']))}",
                f"          value: {defaults['presence_penalty']}",
            ]
        )
    if defaults["reasoning_parsing"] is not None:
        lines.extend(
            [
                "      - key: llm.prediction.reasoning.parsing",
                "        value:",
                f"          enabled: {_bool_text(bool(defaults['reasoning_parsing']['enabled']))}",
                f"          startString: \"{defaults['reasoning_parsing']['startString']}\"",
                f"          endString: \"{defaults['reasoning_parsing']['endString']}\"",
            ]
        )
    if profile["chat_template"]:
        lines.extend(
            [
                "      - key: llm.prediction.promptTemplate",
                "        value:",
                "          type: jinja",
                "          jinjaPromptTemplate:",
                "            template: |",
            ]
        )
        for line in str(profile["chat_template"]).splitlines():
            lines.append(f"              {line}")
        lines.extend(
            [
                "          stopStrings: []",
            ]
        )


def _append_custom_fields(lines: list[str], profile: dict[str, object]) -> None:
    if not profile["needs_enable_thinking"]:
        return
    lines.extend(
        [
            "customFields:",
            "  - key: enableThinking",
            "    displayName: Enable Thinking",
            "    description: Controls whether the model will think before replying",
            "    type: boolean",
            "    defaultValue: true",
            "    effects:",
            "      - type: setJinjaVariable",
            "        variable: enable_thinking",
        ]
    )
    if profile["needs_preserve_thinking"]:
        lines.extend(
            [
                "  - key: preserveThinking",
                "    displayName: Preserve Thinking",
                "    description: Preserve reasoning content in all prior assistant turns instead of only the most recent one",
                "    type: boolean",
                "    defaultValue: false",
                "    effects:",
                "      - type: setJinjaVariable",
                "        variable: preserve_thinking",
            ]
        )


def build_virtual_model_yaml(
    publisher: str,
    slug: str,
    model_key: str,
    profile: dict[str, object],
) -> str:
    normalized_publisher = normalize_owner(publisher)
    lines = [
        "# model.yaml is an open standard for defining cross-platform, composable AI models",
        "# Learn more at https://modelyaml.org",
        f"model: {normalized_publisher}/{slug}",
    ]
    lines.extend(_build_base_block(model_key))
    _append_metadata_overrides(lines, profile)
    _append_operation_fields(lines, profile)
    _append_custom_fields(lines, profile)
    return "\n".join(lines).rstrip() + "\n"


def build_virtual_model_manifest(
    publisher: str,
    slug: str,
    model_key: str,
) -> dict:
    normalized_publisher = normalize_owner(publisher)
    return {
        "type": "model",
        "owner": normalized_publisher,
        "name": slug,
        "dependencies": [
            {
                "type": "model",
                "purpose": "baseModel",
                "modelKeys": [model_key],
                "sources": [],
            }
        ],
        "revision": 1,
    }


def build_readme(publisher: str, slug: str, model_key: str) -> str:
    normalized_publisher = normalize_owner(publisher)
    return (
        f"# {normalized_publisher}/{slug}\n\n"
        "这个目录由 `wrapper_generator.py` 自动生成，用于把本地 concrete GGUF 提升为 LM Studio 可识别的 virtual model。\n\n"
        "## Base Model\n\n"
        f"- `{model_key}`\n"
    )


def build_virtual_model_files(
    publisher: str,
    slug: str,
    model_key: str,
    profile: dict[str, object],
) -> dict[str, str]:
    return {
        "model.yaml": build_virtual_model_yaml(
            publisher=publisher,
            slug=slug,
            model_key=model_key,
            profile=profile,
        ),
        "manifest.json": json.dumps(
            build_virtual_model_manifest(
                publisher=publisher,
                slug=slug,
                model_key=model_key,
            ),
            ensure_ascii=False,
            indent=2,
        )
        + "\n",
        "README.md": build_readme(
            publisher=publisher,
            slug=slug,
            model_key=model_key,
        ),
    }


def scan_gguf_files(models_root: Path) -> list[Path]:
    if not models_root.exists():
        return []
    return sorted(models_root.rglob("*.gguf"))


def should_skip_gguf(models_root: Path, gguf_path: Path) -> str | None:
    relative_parts = [part.lower() for part in gguf_path.relative_to(models_root).parts]
    if "lmstudio-community" in relative_parts:
        return "lmstudio-community"
    if "mmproj" in gguf_path.name.lower():
        return "mmproj"
    return None


def collect_virtual_model_job(
    models_root: Path,
    output_root: Path,
    gguf_path: Path,
    cache_index: dict[str, dict[str, object]],
) -> dict[str, object]:
    skip_reason = should_skip_gguf(models_root, gguf_path)
    raw_publisher = gguf_path.relative_to(models_root).parts[0]
    publisher = normalize_owner(raw_publisher)
    if skip_reason:
        return {
            "source_gguf": str(gguf_path),
            "publisher": publisher,
            "reason": skip_reason,
        }

    metadata = resolve_gguf_metadata(gguf_path, cache_index)
    directory_capabilities = collect_directory_capabilities(gguf_path)
    profile = build_virtual_model_profile(
        metadata,
        has_sibling_mmproj=bool(directory_capabilities["has_sibling_mmproj"]),
    )
    slug = infer_virtual_model_slug(gguf_path.parent.name, metadata)
    model_key = infer_model_key(models_root, gguf_path)
    output_dir = output_root / publisher / slug
    return {
        "source_gguf": str(gguf_path),
        "publisher": publisher,
        "slug": slug,
        "model_key": model_key,
        "output_dir": str(output_dir),
        "capabilities": profile["capabilities"],
        "profile": profile,
        "will_overwrite": any((output_dir / name).exists() for name in ("model.yaml", "manifest.json", "README.md")),
    }


def write_virtual_model(output_dir: Path, publisher: str, slug: str, model_key: str, profile: dict[str, object]) -> None:
    files = build_virtual_model_files(
        publisher=publisher,
        slug=slug,
        model_key=model_key,
        profile=profile,
    )
    output_dir.mkdir(parents=True, exist_ok=True)
    for name, content in files.items():
        (output_dir / name).write_text(content, encoding="utf-8")


def delete_model_index_cache(cache_path: Path | None = None) -> tuple[bool, str | None]:
    if cache_path is None:
        cache_path = DEFAULT_MODEL_INDEX_CACHE
    if not cache_path.exists():
        return False, None
    try:
        cache_path.unlink()
    except OSError as exc:
        return False, str(exc)
    return True, None


def generate_virtual_models_for_directory(
    models_root: Path,
    output_root: Path,
    dry_run: bool,
    cache_path: Path = DEFAULT_GGUF_METADATA_CACHE,
    model_index_cache_path: Path | None = None,
) -> dict[str, object]:
    generated: list[dict[str, object]] = []
    skipped: list[dict[str, object]] = []
    cache_index = load_gguf_metadata_cache(cache_path)

    for gguf_path in scan_gguf_files(models_root):
        job = collect_virtual_model_job(models_root, output_root, gguf_path, cache_index)
        if "reason" in job:
            skipped.append(job)
            continue

        summary_item = {
            "source_gguf": job["source_gguf"],
            "publisher": job["publisher"],
            "slug": job["slug"],
            "model_key": job["model_key"],
            "output_dir": job["output_dir"],
            "capabilities": job["capabilities"],
            "will_overwrite": job["will_overwrite"],
        }
        generated.append(summary_item)

        if not dry_run:
            write_virtual_model(
                output_dir=Path(job["output_dir"]),
                publisher=str(job["publisher"]),
                slug=str(job["slug"]),
                model_key=str(job["model_key"]),
                profile=dict(job["profile"]),
            )

    if model_index_cache_path is None:
        model_index_cache_path = DEFAULT_MODEL_INDEX_CACHE

    deleted_model_index_cache = False
    model_index_cache_delete_error = None
    if not dry_run and generated:
        deleted_model_index_cache, model_index_cache_delete_error = delete_model_index_cache(
            model_index_cache_path
        )

    return {
        "generated_at": datetime.now().isoformat(),
        "dry_run": dry_run,
        "models_root": str(models_root),
        "output_root": str(output_root),
        "deleted_model_index_cache": deleted_model_index_cache,
        "model_index_cache_path": str(model_index_cache_path),
        "model_index_cache_delete_error": model_index_cache_delete_error,
        "generated": generated,
        "skipped": skipped,
    }


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--models-root", default=str(DEFAULT_MODELS_ROOT))
    parser.add_argument("--output-root", default=str(DEFAULT_OUTPUT_ROOT))
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()

    summary = generate_virtual_models_for_directory(
        models_root=Path(args.models_root).expanduser(),
        output_root=Path(args.output_root).expanduser(),
        dry_run=args.dry_run,
    )
    print(json.dumps(summary, ensure_ascii=False, indent=2))
    return 0


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