import asyncio
import importlib
import inspect
import textwrap
from collections import defaultdict
from pathlib import Path
from typing import Any, cast, get_args, get_origin, get_type_hints

import aiohttp_cors
import click
import trafaret as t
from aiohttp import web
from aiohttp.web_urldispatcher import AbstractResource, DynamicResource
from pydantic import BaseModel, TypeAdapter
from trafaret.lib import _empty

import ai.backend.common.validators as tx
from ai.backend.common.api_handlers import BodyParam, PathParam, QueryParam
from ai.backend.common.json import pretty_json_str
from ai.backend.manager import __version__
from ai.backend.manager.api import ManagerStatus
from ai.backend.manager.api.session import UndefChecker
from ai.backend.manager.api.utils import Undefined
from ai.backend.manager.models.vfolder import VFolderPermissionValidator
from ai.backend.manager.server import global_subapp_pkgs


class ParseError(Exception):
    pass


def get_path_parameters(resource: AbstractResource) -> list[dict[str, Any]]:
    params = []
    if isinstance(resource, DynamicResource):
        if groupindex := resource._pattern.groupindex:
            params = [
                {"name": param, "in": "path", "required": True, "schema": {"type": "string"}}
                for param in groupindex.keys()
            ]
    return params


def flatten_or(scheme: t.Trafaret) -> list[t.Trafaret]:
    left, right = scheme.trafarets  # type: ignore[attr-defined]
    items = []
    if isinstance(left, t.Or):
        items.extend(flatten_or(left))
    else:
        items.append(left)

    if isinstance(right, t.Or):
        items.extend(flatten_or(right))
    else:
        items.append(right)
    return items


def _traverse(scheme: t.Trafaret) -> dict[str, Any]:
    if isinstance(scheme, t.Or):
        trafarets = flatten_or(scheme)
        valid_trafarets = [x for x in trafarets if not isinstance(x, (t.Null, UndefChecker))]
        if len(valid_trafarets) >= 2:
            return {"anyOf": [_traverse(s) for s in valid_trafarets]}
        scheme = valid_trafarets[0]
    if isinstance(scheme, t.Any):
        return {"type": "string"}
    if isinstance(scheme, t.Bool):
        return {"type": "boolean"}
    if isinstance(scheme, t.Dict):
        items = parse_trafaret_definition(scheme)
        properties = {d["name"]: d["schema"] for d in items}
        required_keys: list[str] = [d["name"] for d in items if d["required"]]
        return {"type": "object", "properties": properties, "required": required_keys}
    if isinstance(scheme, t.Enum):
        enum_values = scheme.variants  # type: ignore[attr-defined]
        return {"type": "string", "enum": enum_values}
    if isinstance(scheme, tx.DelimiterSeperatedList):
        return {
            "type": "array",
            "items": {"type": "string"},
            "description": (
                "List of items separated by a delimiter (default: comma). "
                "Items may contain spaces, but not the delimiter."
            ),
        }
    if isinstance(scheme, t.Float):
        resp = {"type": "integer"}
        if gte := scheme.gte:  # type: ignore[attr-defined]
            resp["minimum"] = gte
        if lte := scheme.lte:  # type: ignore[attr-defined]
            resp["maximum"] = lte
        return resp
    if isinstance(scheme, t.Int):
        resp = {"type": "integer"}
        if gte := scheme.gte:  # type: ignore[attr-defined]
            resp["minimum"] = gte
        if lte := scheme.lte:  # type: ignore[attr-defined]
            resp["maximum"] = lte
        return resp
    if isinstance(scheme, t.List):
        array_items = _traverse(scheme.trafaret)  # type: ignore[attr-defined]
        return {"type": "array", "items": array_items}
    if isinstance(scheme, t.Mapping):
        key_type = scheme.key.__class__.__name__  # type: ignore[attr-defined]
        value_type = scheme.value.__class__.__name__  # type: ignore[attr-defined]
        return {
            "type": "object",
            "description": f"Mapping({key_type} => {value_type})",
            "additionalProperties": True,
        }
    if isinstance(scheme, t.Regexp):
        pattern = scheme.regexp.pattern  # type: ignore[attr-defined]
        return {"type": "string", "pattern": str(pattern)}
    if isinstance(scheme, t.String):
        return {"type": "string"}
    if isinstance(scheme, t.ToBool):
        return {"type": "boolean"}
    if isinstance(scheme, tx.BinarySize):
        return {"type": "string", "description": "Size in binary format (e.g. 2KB, 3M, 4GiB)"}
    if isinstance(scheme, tx.Enum):
        if scheme.use_name:
            values = [e.name for e in scheme.enum_cls]
        else:
            values = [e.value for e in scheme.enum_cls]
        return {"type": "string", "enum": values}
    if isinstance(scheme, tx.JSONString):
        return {"type": "string", "description": "JSON string"}
    if isinstance(scheme, tx.Path):
        return {"type": "string", "description": "POSIX path"}
    if isinstance(scheme, tx.Slug):
        return {
            "type": "string",
            "description": (
                "String composed of alpha-numeric characters with hyphen and underscores in the middle "
                "(space and dot may be additionally allowed depending on the used locations)"
            ),
        }
    if isinstance(scheme, tx.TimeDuration):
        return {
            "anyOf": [
                {
                    "type": "string",
                    "description": (
                        "Human-readable time duration representation (e.g. 2y, 3d, 4m, 5h, 6s, ...)"
                    ),
                },
                {"type": "int", "minimum": 0, "description": "Number of seconds"},
                {"type": "float", "minimum": 0, "description": "Number of seconds"},
            ]
        }
    if isinstance(scheme, tx.URL):
        return {"type": "string", "format": "uri"}
    if isinstance(scheme, tx.UUID):
        return {"type": "string", "format": "uuid"}
    if isinstance(scheme, VFolderPermissionValidator):
        return {"type": "string", "enum": ["ro", "rw", "wd"]}

    if scheme == t.Email:
        return {"type": "string", "format": "email"}
    if scheme == t.IPv4:
        return {"type": "string", "format": "ipv4"}
    if scheme == t.IPv6:
        return {"type": "string", "format": "ipv6"}
    if scheme == t.URL:
        return {"type": "string", "format": "uri"}

    raise ParseError(f"Failed to convert unknown value {scheme} to OpenAPI type")


def parse_trafaret_value(scheme: t.Trafaret) -> tuple[dict[str, Any], bool]:
    optional = (
        isinstance(scheme, t.Or)
        and len([
            x
            for x in scheme.trafarets  # type: ignore[attr-defined]
            if isinstance(x, (t.Null, UndefChecker))
        ])
        > 0
    )

    return _traverse(scheme), optional


def parse_trafaret_definition(root: t.Dict) -> list[dict[str, Any]]:
    resp = []
    for key in root.keys:  # type: ignore[attr-defined]
        names: list[str] = []
        if isinstance(key, tx.AliasedKey):
            names.extend(key.names)
        elif isinstance(key, t.Key):
            names.append(key.name)  # type: ignore[attr-defined]
        schema, optional = parse_trafaret_value(key.trafaret)
        if key.default and key.default != _empty:
            if inspect.isclass(key.default):
                try:
                    default_value = key.default()
                except Exception:
                    default_value = str(key.default)
            elif isinstance(key.default, Undefined):
                default_value = None
            else:
                default_value = key.default

            schema["default"] = default_value
        if hasattr(key, "__openapi_desc__"):
            schema["description"] = key.__openapi_desc__
        resp += [{"name": names[0], "schema": schema, "required": not optional}]
    return resp


def extract_api_handler_params(
    handler: Any,
) -> tuple[type[BaseModel] | None, type[BaseModel] | None, type[BaseModel] | None]:
    """
    Extract Pydantic models from @api_handler decorated functions.

    Uses get_type_hints() to resolve string annotations from PEP 563
    (`from __future__ import annotations`). inspect.signature() returns
    string annotations in that case, which don't work with get_origin()/get_args().

    Returns:
        tuple of (body_model, query_model, path_model) - any can be None if not present
    """
    body_model: type[BaseModel] | None = None
    query_model: type[BaseModel] | None = None
    path_model: type[BaseModel] | None = None

    try:
        type_hints = get_type_hints(handler)
    except (ValueError, TypeError, NameError):
        return None, None, None

    for param_name, annotation in type_hints.items():
        if param_name == "return":
            continue

        origin = get_origin(annotation)
        if origin is None:
            continue

        args = get_args(annotation)
        if not args:
            continue

        model = args[0]
        if not (isinstance(model, type) and issubclass(model, BaseModel)):
            continue

        if origin is BodyParam:
            body_model = model
        elif origin is QueryParam:
            query_model = model
        elif origin is PathParam:
            path_model = model
        # HeaderParam is typically not included in OpenAPI request body/params

    return body_model, query_model, path_model


def generate_openapi(subapps: list[web.Application], verbose: bool = False) -> dict[str, Any]:
    openapi: dict[str, Any] = {
        "openapi": "3.1.0",
        "info": {
            "title": "Backend.AI Manager API",
            "description": "Backend.AI Manager REST API specification",
            "version": __version__,
            "contact": {
                "name": "Lablup Inc.",
                "url": "https://docs.backend.ai",
                "email": "contect@lablup.com",
            },
        },
        "components": {
            "securitySchemes": {
                "TokenAuth": {
                    "type": "ApiKey",
                    "in": "header",
                    "name": "Authorization: BackendAI",
                    "description": (
                        "Check https://docs.backend.ai/en/latest/manager/common-api/auth.html for"
                        " more information"
                    ),
                },
            },
            "schemas": {},
        },
        "paths": defaultdict(dict),
    }
    operation_id_mapping: defaultdict[str, int] = defaultdict(lambda: 0)
    for app in subapps:
        prefix = app.get("prefix", "root")
        for route in app.router.routes():
            resource = route.resource
            if not resource:
                continue

            if "_root_app" not in app:
                path = "/" + ("" if prefix == "root" else prefix) + resource.canonical
            else:
                path = resource.canonical
            method = route.method

            if method == "OPTIONS":
                continue

            operation_id = f"{prefix}.{route.handler.__name__}"
            if verbose:
                print(f"parsing {operation_id}")
            operation_id_mapping[operation_id] += 1
            if (operation_id_count := operation_id_mapping[operation_id]) > 1:
                operation_id += f".{operation_id_count}"

            description = []
            if route.handler.__doc__:
                description.append(textwrap.dedent(route.handler.__doc__))

            route_def = {
                "operationId": operation_id,
                "tags": [prefix],
                "responses": {"200": {"description": "Successful response"}},
            }
            parameters = []
            parameters.extend(get_path_parameters(resource))
            if hasattr(route.handler, "_backend_attrs"):
                preconds = []
                handler_attrs = route.handler._backend_attrs
                if handler_attrs.get("auth_required"):
                    route_def["security"] = [{"TokenAuth": []}]
                if auth_scope := handler_attrs.get("auth_scope"):
                    preconds.append(f"{auth_scope.capitalize()} privilege required.")
                if manager_status := handler_attrs.get("required_server_statuses"):
                    manager_status = cast(frozenset[ManagerStatus], manager_status)
                    if len(manager_status) == 1:
                        preconds.append(
                            f"Manager status required: {list(manager_status)[0].value.upper()}"
                        )
                    else:
                        preconds.append(
                            "Manager status required: one of "
                            f"{', '.join([e.value.upper() for e in sorted(manager_status)])}"
                        )
                if preconds:
                    description.append("\n**Preconditions:**")
                    for item in preconds:
                        description.append(f"* {item}")
                    description.append("")
                if request_scheme := handler_attrs.get("request_scheme"):
                    if isinstance(request_scheme, t.Dict):
                        parsed_definition = parse_trafaret_definition(request_scheme)
                        if method == "GET" or method == "DELETE":
                            parameters.extend([{**d, "in": "query"} for d in parsed_definition])
                        else:
                            properties = {d["name"]: d["schema"] for d in parsed_definition}
                            required_keys: list[str] = [
                                d["name"] for d in parsed_definition if d["required"]
                            ]
                            raw_examples = handler_attrs.get("request_examples") or []
                            examples = {
                                f"{operation_id}_Example{i}": {"value": e}
                                for e, i in zip(
                                    raw_examples, range(1, len(raw_examples) + 1), strict=True
                                )
                            }
                            route_def["requestBody"] = {
                                "content": {
                                    "application/json": {
                                        "schema": {
                                            "type": "object",
                                            "properties": properties,
                                            "required": required_keys,
                                        },
                                        "examples": examples,
                                    }
                                }
                            }
                    elif issubclass(request_scheme, BaseModel):
                        schema_name = request_scheme.__name__
                        request_schema = request_scheme.model_json_schema(
                            ref_template="#/components/schemas/{model}"
                        )

                        if additional_definitions := request_schema.pop("$defs", None):
                            openapi["components"]["schemas"].update(additional_definitions)
                        openapi["components"]["schemas"][schema_name] = request_schema
                        route_def["requestBody"] = {
                            "content": {
                                "application/json": {
                                    "schema": {"$ref": f"#/components/schemas/{schema_name}"}
                                }
                            }
                        }
                    else:
                        raise RuntimeError(
                            f"{request_scheme} not considered as a valid request type"
                        )

            # Handle @api_handler decorated functions (BodyParam, QueryParam, PathParam)
            if "requestBody" not in route_def:
                body_model, query_model, path_model = extract_api_handler_params(route.handler)

                if body_model is not None:
                    schema_name = body_model.__name__
                    request_schema = body_model.model_json_schema(
                        ref_template="#/components/schemas/{model}"
                    )
                    if additional_definitions := request_schema.pop("$defs", None):
                        openapi["components"]["schemas"].update(additional_definitions)
                    openapi["components"]["schemas"][schema_name] = request_schema
                    route_def["requestBody"] = {
                        "content": {
                            "application/json": {
                                "schema": {"$ref": f"#/components/schemas/{schema_name}"}
                            }
                        }
                    }

                if query_model is not None:
                    query_schema = query_model.model_json_schema(
                        ref_template="#/components/schemas/{model}"
                    )
                    properties = query_schema.get("properties", {})
                    required_keys = query_schema.get("required", [])
                    for prop_name, prop_schema in properties.items():
                        parameters.append({
                            "name": prop_name,
                            "in": "query",
                            "required": prop_name in required_keys,
                            "schema": prop_schema,
                        })

                if path_model is not None:
                    path_schema = path_model.model_json_schema(
                        ref_template="#/components/schemas/{model}"
                    )
                    properties = path_schema.get("properties", {})
                    for prop_name, prop_schema in properties.items():
                        # Check if already added by get_path_parameters
                        if not any(p["name"] == prop_name for p in parameters):
                            parameters.append({
                                "name": prop_name,
                                "in": "path",
                                "required": True,
                                "schema": prop_schema,
                            })

                # TODO: Extract response model from @api_handler decorated functions.
                # Currently, @api_handler returns APIResponse which wraps BaseResponseModel,
                # but the actual response model type is not captured in the return type annotation.
                #
                # To enable OpenAPI response schema generation, APIResponse should be made Generic:
                #
                #   # In api_handlers.py:
                #   TResponseModel = TypeVar("TResponseModel", bound=BaseResponseModel)
                #
                #   @dataclass
                #   class APIResponse(Generic[TResponseModel]):
                #       _status_code: int
                #       _data: Optional[TResponseModel]
                #
                #   # In handler:
                #   async def handler(...) -> APIResponse[MyResponseModel]:
                #       return APIResponse.build(200, MyResponseModel(...))
                #
                # Then we can extract the response model via:
                #   ret_type = get_type_hints(handler).get("return")
                #   if get_origin(ret_type) is APIResponse:
                #       response_model = get_args(ret_type)[0]

            route_def["parameters"] = parameters
            route_def["description"] = "\n".join(description)

            # Extract response schema from handler's return type annotation.
            # Supported return types:
            #   - BaseModel subclass: async def handler(...) -> MyResponseModel
            #   - list[BaseModel]: async def handler(...) -> list[MyResponseModel]
            #
            # Note: @api_handler decorated functions return APIResponse (not BaseModel),
            # so they are not captured here. See TODO comment above for future improvement.
            try:
                type_hints = get_type_hints(route.handler)
            except (NameError, AttributeError):
                # Skip type hint extraction for handlers with unresolvable forward references
                # (e.g., GraphQLView subclasses)
                type_hints = {}
            if (
                (ret_type := type_hints.get("return"))
                # For generic types like list[Model], get the origin (list).
                # For non-generic types like Model, use the type itself.
                and (response_cls := getattr(ret_type, "__origin__", ret_type))
                and (issubclass(response_cls, BaseModel) or issubclass(response_cls, list))
            ):
                response_schema: dict[str, Any]
                if issubclass(response_cls, list):
                    # Handle list[SomeModel] return type
                    arg: type[BaseModel]
                    (arg,) = get_args(ret_type)
                    schema_name = f"{arg.__name__}_List"
                    response_schema = TypeAdapter(list[arg]).json_schema(  # type: ignore[valid-type]
                        ref_template="#/components/schemas/{model}"
                    )
                elif issubclass(response_cls, BaseModel):
                    # Handle direct BaseModel return type
                    schema_name = response_cls.__name__
                    response_schema = response_cls.model_json_schema(
                        ref_template="#/components/schemas/{model}"
                    )

                else:
                    raise RuntimeError(f"{response_cls} not considered as a valid response type")

                # Extract nested model definitions (e.g., referenced models) to components
                if additional_definitions := response_schema.pop("$defs", None):
                    openapi["components"]["schemas"].update(additional_definitions)
                openapi["components"]["schemas"][schema_name] = response_schema
                route_def["responses"] = {
                    "200": {
                        "description": "",
                        "content": {
                            "application/json": {
                                "schema": {"$ref": f"#/components/schemas/{schema_name}"}
                            }
                        },
                    }
                }
            openapi["paths"][path][method.lower()] = route_def
    return openapi


async def generate() -> dict[str, Any]:
    cors_options = {
        "*": aiohttp_cors.ResourceOptions(  # type: ignore[no-untyped-call]
            allow_credentials=False, expose_headers="*", allow_headers="*"
        ),
    }

    subapps: list[web.Application] = []
    for subapp in global_subapp_pkgs:
        pkg = importlib.import_module("ai.backend.manager.api" + subapp)
        app, _ = pkg.create_app(cors_options)
        subapps.append(app)
    return generate_openapi(subapps, verbose=True)


@click.command()
@click.option(
    "--output",
    "-o",
    default=None,
    type=click.Path(dir_okay=False, writable=True),
    help="Output file path (default: stdout)",
)
def main(output: Path | None) -> None:
    """
    Generates OpenAPI specification of Backend.AI API.
    """
    openapi = asyncio.run(generate())
    if output is None:
        print(pretty_json_str(openapi))
    else:
        with output.open(mode="w") as fw:
            fw.write(pretty_json_str(openapi))


if __name__ == "__main__":
    main()
