"""
Per-feature OpenAPI snapshot for lazy-loaded routers.

The committed JSON is generated by `python -m litellm.proxy._lazy_openapi_snapshot`
and consumed at runtime so /openapi.json can show full route info for unloaded
features without importing them. CI verifies the file is current and surfaces
any drift as a neutral check.
"""

import json
import re
import sys
from pathlib import Path
from typing import Dict, Optional, Set

SNAPSHOT_FILE = Path(__file__).parent / "_lazy_openapi_snapshot.json"
HTTP_METHOD_SUFFIXES = {
    "delete",
    "get",
    "head",
    "options",
    "patch",
    "post",
    "put",
    "trace",
}


def _stabilize_multi_method_route_ids(routes) -> None:
    """FastAPI derives route IDs from a set of methods; make snapshots stable."""

    for route in routes:
        methods = sorted(getattr(route, "methods", None) or [])
        if len(methods) <= 1 or not getattr(route, "path_format", None):
            continue

        operation_id = f"{route.name}{route.path_format}"
        operation_id = re.sub(r"\W", "_", operation_id)
        route.unique_id = f"{operation_id}_{methods[0].lower()}"


def load_snapshot() -> Optional[Dict[str, Dict]]:
    if not SNAPSHOT_FILE.exists():
        return None
    try:
        with SNAPSHOT_FILE.open() as f:
            return json.load(f)
    except (json.JSONDecodeError, OSError):
        return None


def _normalize_operation_ids(paths: Dict[str, Dict]) -> None:
    """Make FastAPI-generated operation IDs stable for multi-method routes.

    FastAPI derives the default operation ID suffix from the first item in the
    route's methods set. For routes registered with several HTTP methods, that
    set iteration order can vary between processes, which makes the snapshot
    drift even when no routes changed.
    """
    for path_ops in paths.values():
        if not isinstance(path_ops, dict):
            continue

        methods = {method for method in path_ops if method in HTTP_METHOD_SUFFIXES}
        if not methods:
            continue

        for method, operation in path_ops.items():
            if method not in HTTP_METHOD_SUFFIXES or not isinstance(operation, dict):
                continue

            operation_id = operation.get("operationId")
            if not isinstance(operation_id, str):
                continue

            for suffix in methods:
                suffix_token = f"_{suffix}"
                if operation_id.endswith(suffix_token):
                    operation["operationId"] = (
                        operation_id[: -len(suffix_token)] + f"_{method}"
                    )
                    break


def generate_snapshot() -> Dict[str, Dict]:
    import importlib

    from fastapi.openapi.utils import get_openapi

    from litellm.proxy._lazy_features import LAZY_FEATURES
    from litellm.proxy.proxy_server import app, ensure_unique_openapi_operation_ids

    for feat in LAZY_FEATURES:
        if feat.module_path in sys.modules:
            continue
        try:
            module = importlib.import_module(feat.module_path)
            feat.register_fn(app, module)
        except Exception as exc:
            sys.stderr.write(f"warning: skip {feat.name}: {exc}\n")

    fragments: Dict[str, Dict] = {}
    used_operation_ids: Set[str] = set()
    for feat in LAZY_FEATURES:
        feat_routes = [
            r
            for r in app.routes
            if any(getattr(r, "path", "").startswith(p) for p in feat.path_prefixes)
        ]
        if not feat_routes:
            continue
        _stabilize_multi_method_route_ids(feat_routes)
        full = get_openapi(title=app.title, version=app.version, routes=feat_routes)
        paths = full.get("paths", {})
        _normalize_operation_ids(paths)
        # Group all of a feature's routes under one tag.
        for path_ops in full.get("paths", {}).values():
            for method, op in path_ops.items():
                if isinstance(op, dict):
                    operation_id = op.get("operationId")
                    if isinstance(operation_id, str):
                        for suffix in HTTP_METHOD_SUFFIXES:
                            if operation_id.endswith(f"_{suffix}"):
                                op["operationId"] = (
                                    operation_id[: -len(suffix)] + method
                                )
                                break
                    op["tags"] = [feat.name]
        full = ensure_unique_openapi_operation_ids(full, used_operation_ids)
        fragments[feat.name] = {
            "paths": paths,
            "components": {"schemas": full.get("components", {}).get("schemas", {})},
        }
    return fragments


if __name__ == "__main__":
    fragments = generate_snapshot()
    SNAPSHOT_FILE.write_text(json.dumps(fragments, indent=2, sort_keys=True) + "\n")
    sys.stdout.write(f"wrote {len(fragments)} feature fragments to {SNAPSHOT_FILE}\n")