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

DEFAULT_OFFICIAL_MODEL_YAML = Path(r"~\hub\models\qwen\qwen3.6-35b-a3b\model.yaml")
DEFAULT_OFFICIAL_MANIFEST = Path(r"~\hub\models\qwen\qwen3.6-35b-a3b\manifest.json")
DEFAULT_OUTPUT_ROOT = Path("artifacts/lmstudio-model-wrappers")
DEFAULT_TARGET_MODEL_KEY = "qwen3.6-35b-a3b-uncensored-hauhaucs-aggressive"
DEFAULT_WRAPPER_OWNER = "local"
DEFAULT_WRAPPER_NAME = "qwen3.6-35b-a3b-hauhaucs-aggressive"
DEFAULT_MODELS_ROOT = Path.home() / ".lmstudio" / "models"
DEFAULT_HUB_ROOT = Path.home() / ".lmstudio" / "hub" / "models"


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 _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 read_gguf_metadata(path: Path) -> dict[str, object]:
    data = path.read_bytes()
    if len(data) < 24 or data[:4] != b"GGUF":
        raise ValueError(f"Not a GGUF file: {path}")

    version, offset = _read_struct(data, 4, "<I")
    if int(version) not in (2, 3):
        raise ValueError(f"Unsupported GGUF version: {version}")

    _, offset = _read_struct(data, offset, "<Q")
    metadata_count, offset = _read_struct(data, offset, "<Q")

    metadata: dict[str, object] = {}
    for _ in range(int(metadata_count)):
        key, offset = _read_gguf_string(data, offset)
        value_type, offset = _read_struct(data, offset, "<I")
        value, offset = _read_gguf_value(data, offset, int(value_type))
        metadata[key] = value
    return metadata


def extract_architecture_family(value: object) -> str:
    text = str(value or "").strip().lower()
    if not text:
        return ""
    head = []
    for char in text:
        if char.isalpha():
            head.append(char)
            continue
        if head:
            break
    return "".join(head) or text


def infer_model_identity(
    models_root: Path, gguf_path: Path, metadata: dict[str, object]
) -> tuple[str, str]:
    try:
        relative = gguf_path.relative_to(models_root)
        parts = list(relative.parts[:-1])
    except ValueError:
        parts = []

    owner = ""
    name = ""
    if len(parts) >= 2:
        owner = parts[0]
        name = parts[1]
    elif len(parts) == 1:
        owner = "local"
        name = parts[0]

    if not owner:
        owner = "local"

    if not name:
        name = str(metadata.get("general.name") or gguf_path.stem)

    owner = owner.strip().replace("\\", "-").replace("/", "-")
    name = name.strip().replace("\\", "-").replace("/", "-")
    return owner, name


def infer_model_key(gguf_path: Path) -> str:
    return gguf_path.stem


def _metadata_bool(metadata: dict[str, object], keys: list[str]) -> bool | None:
    for key in keys:
        if key in metadata and isinstance(metadata[key], bool):
            return bool(metadata[key])
    return None


def extract_capabilities_from_metadata(
    metadata: dict[str, object],
) -> dict[str, bool | None]:
    reasoning = _metadata_bool(
        metadata,
        [
            "capabilities.reasoning",
            "general.reasoning",
            "reasoning",
        ],
    )
    if reasoning is None:
        chat_template = metadata.get("tokenizer.chat_template")
        if isinstance(chat_template, str) and "enable_thinking" in chat_template:
            reasoning = True

    vision = _metadata_bool(
        metadata,
        [
            "capabilities.vision",
            "general.vision",
            "vision",
        ],
    )
    tools = _metadata_bool(
        metadata,
        [
            "capabilities.tools",
            "general.tools",
            "tools",
        ],
    )
    return {
        "reasoning": reasoning,
        "vision": vision,
        "tools": tools,
    }


def extract_capabilities_from_template(model_yaml_text: str) -> dict[str, bool | None]:
    lowered = model_yaml_text.lower()
    reasoning = None
    if "reasoning: true" in lowered or "key: enablethinking" in lowered:
        reasoning = True
    elif "reasoning: false" in lowered:
        reasoning = False

    vision = True if "vision: true" in lowered else None
    tools = True if "tools: true" in lowered else None
    return {
        "reasoning": reasoning,
        "vision": vision,
        "tools": tools,
    }


def merge_capabilities(
    metadata_capabilities: dict[str, bool | None],
    template_capabilities: dict[str, bool | None],
) -> dict[str, bool | None]:
    merged: dict[str, bool | None] = {}
    for key in ("reasoning", "vision", "tools"):
        merged[key] = (
            metadata_capabilities[key]
            if metadata_capabilities.get(key) is not None
            else template_capabilities.get(key)
        )
    return merged


def find_template_package(hub_root: Path, metadata: dict[str, object]) -> Path | None:
    family = extract_architecture_family(metadata.get("general.architecture"))
    if not family or not hub_root.exists():
        return None

    best_match: tuple[int, Path] | None = None
    for model_yaml_path in hub_root.rglob("model.yaml"):
        package_dir = model_yaml_path.parent
        manifest_path = package_dir / "manifest.json"
        if not manifest_path.exists():
            continue
        try:
            relative = package_dir.relative_to(hub_root)
        except ValueError:
            continue
        if not relative.parts:
            continue

        owner = relative.parts[0].lower()
        name = relative.parts[-1].lower()
        if owner == "lmstudio-community":
            continue

        score = 0
        if owner == family:
            score += 100
        elif family in owner:
            score += 70
        if family in name:
            score += 40

        if score <= 0:
            continue
        if best_match is None or score > best_match[0]:
            best_match = (score, package_dir)

    if best_match is None:
        return None
    return best_match[1]


def _find_line_index(lines: list[str], prefix: str) -> int | None:
    for index, line in enumerate(lines):
        if line.strip().startswith(prefix):
            return index
    return None


def _find_section_bounds(lines: list[str], header: str) -> tuple[int, int] | None:
    for index, line in enumerate(lines):
        if line.strip() != header:
            continue
        indent = len(line) - len(line.lstrip(" "))
        end = len(lines)
        for probe in range(index + 1, len(lines)):
            candidate = lines[probe]
            stripped = candidate.strip()
            if not stripped:
                continue
            candidate_indent = len(candidate) - len(candidate.lstrip(" "))
            if candidate_indent <= indent:
                end = probe
                break
        return index, end
    return None


def _replace_base_block(lines: list[str], target_model_key: str) -> list[str]:
    result: list[str] = []
    index = 0
    replaced = False
    while index < len(lines):
        line = lines[index]
        if not replaced and line.strip().startswith("base:"):
            base_indent = len(line) - len(line.lstrip(" "))
            result.append(f'base: "{target_model_key}"')
            replaced = True
            index += 1
            while index < len(lines):
                candidate = lines[index]
                stripped = candidate.strip()
                if not stripped:
                    index += 1
                    continue
                candidate_indent = len(candidate) - len(candidate.lstrip(" "))
                if candidate_indent <= base_indent:
                    break
                index += 1
            continue
        result.append(line)
        index += 1
    if not replaced:
        raise ValueError("Official model.yaml is missing a base declaration")
    return result


def _upsert_metadata_bool(lines: list[str], key: str, value: bool | None) -> list[str]:
    if value is None:
        return lines

    bounds = _find_section_bounds(lines, "metadataOverrides:")
    if bounds is None:
        insert_at = _find_line_index(lines, "config:")
        metadata_block = [
            "metadataOverrides:",
            f"  {key}: {_bool_text(value)}",
        ]
        if insert_at is None:
            return lines + metadata_block
        return lines[:insert_at] + metadata_block + lines[insert_at:]

    start, end = bounds
    child_prefix = "  "
    for index in range(start + 1, end):
        if lines[index].strip().startswith(f"{key}:"):
            lines[index] = f"{child_prefix}{key}: {_bool_text(value)}"
            return lines

    return lines[:end] + [f"{child_prefix}{key}: {_bool_text(value)}"] + lines[end:]


def build_wrapper_model_yaml(
    official_model_yaml_text: str,
    target_model_key: str,
    wrapper_owner: str,
    wrapper_name: str,
    capabilities: dict[str, bool | None] | None = None,
) -> str:
    lines = official_model_yaml_text.splitlines()
    model_index = _find_line_index(lines, "model:")
    if model_index is None:
        raise ValueError("Official model.yaml is missing a model declaration")
    lines[model_index] = f"model: {wrapper_owner}/{wrapper_name}"

    lines = _replace_base_block(lines, target_model_key)

    if capabilities:
        for key in ("reasoning", "vision", "tools"):
            lines = _upsert_metadata_bool(lines, key, capabilities.get(key))

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


def build_wrapper_manifest(
    official_manifest: dict,
    target_model_key: str,
    wrapper_owner: str,
    wrapper_name: str,
) -> dict:
    manifest = {
        "type": "model",
        "owner": wrapper_owner,
        "name": wrapper_name,
        "dependencies": [
            {
                "type": "model",
                "purpose": "baseModel",
                "modelKeys": [target_model_key],
            }
        ],
        "revision": int(official_manifest.get("revision", 1)),
    }

    source_dep = (official_manifest.get("dependencies") or [{}])[0]
    if source_dep.get("sources"):
        manifest["dependencies"][0]["sources"] = source_dep["sources"]
    return manifest


def build_readme(
    target_model_key: str,
    wrapper_owner: str,
    wrapper_name: str,
) -> str:
    return (
        f"# {wrapper_owner}/{wrapper_name}\n\n"
        "这个目录由 `wrapper_generator.py` 自动生成，用于补全 LM Studio 可识别的模型元数据。\n\n"
        "## 源模型\n\n"
        f"- 本地模型 key: `{target_model_key}`\n\n"
        "## 说明\n\n"
        "- 不修改 GGUF 文件本身\n"
        "- 复用 LM Studio 官方模型包结构\n"
        "- 用于补全模型能力标注与基础依赖关系\n"
    )


def build_wrapper_files(
    target_model_key: str,
    wrapper_owner: str,
    wrapper_name: str,
    official_model_yaml_text: str,
    official_manifest: dict,
    capabilities: dict[str, bool | None] | None = None,
) -> dict:
    model_yaml = build_wrapper_model_yaml(
        official_model_yaml_text=official_model_yaml_text,
        target_model_key=target_model_key,
        wrapper_owner=wrapper_owner,
        wrapper_name=wrapper_name,
        capabilities=capabilities,
    )
    manifest = build_wrapper_manifest(
        official_manifest=official_manifest,
        target_model_key=target_model_key,
        wrapper_owner=wrapper_owner,
        wrapper_name=wrapper_name,
    )
    readme = build_readme(
        target_model_key=target_model_key,
        wrapper_owner=wrapper_owner,
        wrapper_name=wrapper_name,
    )
    return {
        "model.yaml": model_yaml,
        "manifest.json": json.dumps(manifest, ensure_ascii=False, indent=2) + "\n",
        "README.md": readme,
    }


def generate_wrapper(
    output_dir: Path,
    target_model_key: str,
    wrapper_owner: str,
    wrapper_name: str,
    official_model_yaml_text: str,
    official_manifest: dict,
    dry_run: bool,
    capabilities: dict[str, bool | None] | None = None,
) -> dict:
    files = build_wrapper_files(
        target_model_key=target_model_key,
        wrapper_owner=wrapper_owner,
        wrapper_name=wrapper_name,
        official_model_yaml_text=official_model_yaml_text,
        official_manifest=official_manifest,
        capabilities=capabilities,
    )

    summary = {
        "generated_at": datetime.now().isoformat(),
        "dry_run": dry_run,
        "output_dir": str(output_dir),
        "wrapper_model": f"{wrapper_owner}/{wrapper_name}",
        "target_model_key": target_model_key,
        "files": sorted(files.keys()),
    }

    if dry_run:
        return summary

    output_dir.mkdir(parents=True, exist_ok=True)
    for name, content in files.items():
        (output_dir / name).write_text(content, encoding="utf-8")
    (output_dir / "summary.json").write_text(
        json.dumps(summary, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )
    return summary


def load_inputs(model_yaml_path: Path, manifest_path: Path) -> tuple[str, dict]:
    return (
        model_yaml_path.read_text(encoding="utf-8"),
        json.loads(manifest_path.read_text(encoding="utf-8")),
    )


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


def _skip_reason_for_model(
    models_root: Path, gguf_path: Path, owner: str
) -> str | None:
    try:
        relative_parts = {
            part.lower() for part in gguf_path.relative_to(models_root).parts
        }
    except ValueError:
        relative_parts = {part.lower() for part in gguf_path.parts}
    if "lmstudio-community" in relative_parts or owner.lower() == "lmstudio-community":
        return "lmstudio-community"
    return None


def collect_model_job(
    models_root: Path,
    hub_root: Path,
    output_root: Path,
    gguf_path: Path,
) -> dict:
    metadata = read_gguf_metadata(gguf_path)
    owner, name = infer_model_identity(models_root, gguf_path, metadata)
    skip_reason = _skip_reason_for_model(models_root, gguf_path, owner)
    if skip_reason:
        return {
            "source_gguf": str(gguf_path),
            "owner": owner,
            "name": name,
            "reason": skip_reason,
        }

    template_dir = find_template_package(hub_root, metadata)
    if template_dir is None:
        return {
            "source_gguf": str(gguf_path),
            "owner": owner,
            "name": name,
            "reason": "template-not-found",
        }

    output_dir = output_root / owner / name
    if output_dir == template_dir:
        return {
            "source_gguf": str(gguf_path),
            "owner": owner,
            "name": name,
            "reason": "output-conflicts-with-template",
        }

    official_model_yaml_text, official_manifest = load_inputs(
        template_dir / "model.yaml",
        template_dir / "manifest.json",
    )
    metadata_capabilities = extract_capabilities_from_metadata(metadata)
    template_capabilities = extract_capabilities_from_template(official_model_yaml_text)
    capabilities = merge_capabilities(metadata_capabilities, template_capabilities)

    model_key = infer_model_key(gguf_path)
    return {
        "source_gguf": str(gguf_path),
        "owner": owner,
        "name": name,
        "model_key": model_key,
        "template_path": str(template_dir),
        "output_dir": str(output_dir),
        "will_overwrite": any(
            (output_dir / file_name).exists()
            for file_name in ("model.yaml", "manifest.json", "README.md")
        ),
        "capabilities": capabilities,
        "official_model_yaml_text": official_model_yaml_text,
        "official_manifest": official_manifest,
    }


def generate_wrappers_for_models(
    models_root: Path,
    hub_root: Path,
    output_root: Path,
    dry_run: bool,
) -> dict:
    generated = []
    skipped = []

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

        output_dir = Path(job["output_dir"])
        summary_item = {
            "source_gguf": job["source_gguf"],
            "owner": job["owner"],
            "name": job["name"],
            "model_key": job["model_key"],
            "template_path": job["template_path"],
            "output_dir": job["output_dir"],
            "will_overwrite": job["will_overwrite"],
            "capabilities": job["capabilities"],
        }

        if not dry_run:
            generate_wrapper(
                output_dir=output_dir,
                target_model_key=job["model_key"],
                wrapper_owner=job["owner"],
                wrapper_name=job["name"],
                official_model_yaml_text=job["official_model_yaml_text"],
                official_manifest=job["official_manifest"],
                dry_run=False,
                capabilities=job["capabilities"],
            )

        generated.append(summary_item)

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


def _requested_legacy_single_mode(args: argparse.Namespace) -> bool:
    return any(
        [
            args.target_model_key,
            args.wrapper_owner,
            args.wrapper_name,
            args.official_model_yaml,
            args.official_manifest,
            args.output_dir,
        ]
    )


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

    parser.add_argument("--target-model-key", default="")
    parser.add_argument("--wrapper-owner", default="")
    parser.add_argument("--wrapper-name", default="")
    parser.add_argument("--official-model-yaml", default="")
    parser.add_argument("--official-manifest", default="")
    parser.add_argument("--output-dir", default="")
    args = parser.parse_args()

    if _requested_legacy_single_mode(args):
        target_model_key = args.target_model_key or DEFAULT_TARGET_MODEL_KEY
        wrapper_owner = args.wrapper_owner or DEFAULT_WRAPPER_OWNER
        wrapper_name = args.wrapper_name or DEFAULT_WRAPPER_NAME
        official_model_yaml_path = Path(
            args.official_model_yaml or DEFAULT_OFFICIAL_MODEL_YAML
        )
        official_manifest_path = Path(
            args.official_manifest or DEFAULT_OFFICIAL_MANIFEST
        )
        official_model_yaml_text, official_manifest = load_inputs(
            official_model_yaml_path,
            official_manifest_path,
        )
        if args.output_dir:
            output_dir = Path(args.output_dir)
        else:
            output_dir = (
                Path(args.output_root or DEFAULT_OUTPUT_ROOT)
                / wrapper_owner
                / wrapper_name
            )
        summary = generate_wrapper(
            output_dir=output_dir,
            target_model_key=target_model_key,
            wrapper_owner=wrapper_owner,
            wrapper_name=wrapper_name,
            official_model_yaml_text=official_model_yaml_text,
            official_manifest=official_manifest,
            dry_run=args.dry_run,
        )
    else:
        summary = generate_wrappers_for_models(
            models_root=Path(args.models_root).expanduser(),
            hub_root=Path(args.hub_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())