"""
python -m cog.command.openapi_schema

This prints a JSON object describing the OpenAPI schema of the model.
Schema is generated by introspecting the predictor's type annotations
without starting the HTTP server.
"""

import importlib.util
import json
import os
import sys
from typing import Any, Dict

from .._inspector import create_predictor
from .._schemas import to_json_schema
from ..config import Config
from ..errors import ConfigDoesNotExist
from ..mode import Mode
from ..suppress_output import suppress_output


def _load_module_from_ref(ref: str) -> tuple[str, str]:
    """Load a predictor module from a ref like 'predict.py:Predictor' or 'my-subdir/predict.py:Predictor'.

    Uses spec_from_file_location to load the module by file path, which handles
    subdirectory predictors correctly (unlike import_module which requires the
    module to be on sys.path).

    Returns (module_name, class_name) with the module pre-loaded in sys.modules.
    """
    module_path, class_name = ref.rsplit(":", 1) if ":" in ref else (ref, "Predictor")
    module_name = os.path.basename(module_path).removesuffix(".py")

    # Load module from file path so subdirectory predictors work
    spec = importlib.util.spec_from_file_location(module_name, module_path)
    if spec is not None and spec.loader is not None:
        module = importlib.util.module_from_spec(spec)
        sys.modules[module_name] = module
        spec.loader.exec_module(module)

    return module_name, class_name


def remove_title_next_to_ref(
    schema_node: Any,
) -> Any:
    """
    Recursively remove 'title' from schema components that have a '$ref'.
    This function addresses a non-compliance issue in FastAPI's OpenAPI schema generation.
    """
    if isinstance(schema_node, dict):
        if "$ref" in schema_node and "title" in schema_node:
            del schema_node["title"]
        for _key, value in schema_node.items():
            remove_title_next_to_ref(value)
    elif isinstance(schema_node, list):
        for i, item in enumerate(schema_node):
            schema_node[i] = remove_title_next_to_ref(item)
    return schema_node


def fix_nullable_anyof(schema_node: Any) -> None:
    """
    Convert anyOf with null type to nullable: true for OpenAPI 3.0 compatibility.

    FastAPI generates: {"anyOf": [{"type": "string"}, {"type": "null"}]}
    OpenAPI 3.0 wants: {"type": "string", "nullable": true}
    """
    if isinstance(schema_node, dict):
        if "anyOf" in schema_node:
            anyof = schema_node["anyOf"]
            if isinstance(anyof, list) and len(anyof) == 2:
                # Check if one is {"type": "null"}
                null_idx = None
                other_idx = None
                for i, item in enumerate(anyof):
                    if isinstance(item, dict) and item.get("type") == "null":
                        null_idx = i
                    else:
                        other_idx = i

                if null_idx is not None and other_idx is not None:
                    other = anyof[other_idx]
                    if isinstance(other, dict):
                        # Replace anyOf with the non-null type + nullable
                        del schema_node["anyOf"]
                        schema_node.update(other)
                        schema_node["nullable"] = True

        for value in schema_node.values():
            fix_nullable_anyof(value)
    elif isinstance(schema_node, list):
        for item in schema_node:
            fix_nullable_anyof(item)


if __name__ == "__main__":
    schema: Dict[str, Any] = {}
    try:
        config = Config()
        # Determine mode: prefer predict, fall back to train
        try:
            ref = config.get_predictor_ref(Mode.PREDICT)
            mode = Mode.PREDICT
        except ValueError:
            ref = config.get_predictor_ref(Mode.TRAIN)
            mode = Mode.TRAIN

        module_name, class_name = _load_module_from_ref(ref)

        with suppress_output():
            predictor_info = create_predictor(module_name, class_name)

        schema = to_json_schema(predictor_info, mode)
        remove_title_next_to_ref(schema)
        fix_nullable_anyof(schema)
    except FileNotFoundError:
        raise ConfigDoesNotExist("cog.yaml not found") from None

    print(json.dumps(schema, indent=2))
