#!/usr/bin/env python3
import json
import logging
import os
import subprocess  # nosec
from enum import Enum
from itertools import chain
from pathlib import Path
from tempfile import TemporaryDirectory
from typing import (
    Any,
    Callable,
    Dict,
    ItemsView,
    List,
    Literal,
    Optional,
    Tuple,
    TypeGuard,
    Union,
    cast,
)

import openapi_pydantic
import structlog
import typer
from jinja2 import Environment, FileSystemLoader, select_autoescape
from openapi_pydantic import OpenAPI, PathItem, Reference, Schema
from openapi_pydantic import Operation as OpenAPIOperation
from openapi_pydantic import Parameter as OpenApiParameter
from pydantic import BaseModel, ConfigDict

logger = structlog.get_logger(__name__)


class SchemaNotFound(Exception):
    pass


TYPES_MAPPING = {
    "boolean": "bool",
    "integer": "int",
    "string": "str",
}

template_env = Environment(
    loader=FileSystemLoader(Path(__file__).parent / "templates"),
    trim_blocks=True,
    lstrip_blocks=True,
    extensions=["jinja2.ext.loopcontrols"],
    autoescape=select_autoescape(),
)

METHODS = ["get", "put", "post", "patch", "delete"]


def to_python_type(t: str) -> str:
    return TYPES_MAPPING[t]


def to_pascal_case(text: str) -> str:
    """Convert a string to PascalCase.

    Example: 'api_keys' -> 'ApiKeys'
    """
    if len(text) == 0:
        return text

    import string
    return string.capwords(text.replace("_", " ").replace("-", " ")).replace(" ", "")


# Register the function as a Jinja filter
template_env.filters["to_pascal_case"] = to_pascal_case


class ParameterKind(Enum):
    QUERY = "query"
    PATH = "path"
    """When the parameter appears as multipart/form-data"""
    FORM = "form"


class Parameter(BaseModel):
    """
    A parameter of an operation that will get rendered as an argument of a function.

    Attributes
    ----------
        name
            The name of the parameter.
        type
            The type of the parameter.
        required
            Indicates whether the parameter is required.
        kind
            The kind of the parameter.
        is_enum
            Indicates whether the parameter is an enum. Defaults to False.
            If it's true, the type is the name of the enum.
        is_python_type
            Indicates whether the type is already a Python type (e.g., List[str]).
            If True, the type will not be converted.

    """

    name: str
    type: str
    required: bool
    kind: ParameterKind
    is_enum: bool = False
    is_python_type: bool = False

    def to_argument(self) -> str:
        """
        Convert the parameter to an argument string.

        Returns
        -------
            arg: The argument string representation of the parameter.
        """
        # If it's an enumeration or already a Python type, use the type as-is
        if self.is_enum or self.is_python_type:
            non_optional_type = self.type
        else:
            non_optional_type = to_python_type(self.type)

        if self.required:
            arg = f"{self.name}: {non_optional_type},"
        else:
            arg = f"{self.name}: Optional[{non_optional_type}] = None,"
        return arg


class Response(BaseModel):
    ref: str  # e.g. "#/components/schemas/Dataset"
    name: str  # e.g. "Dataset"
    type: Optional[Literal["array"]] = None

    def type_hint(self) -> str:
        if not self.type:
            return self.name
        if self.type == "array":
            return f"List[{self.name}]"
        raise NotImplementedError(f"Type hint for {self.type} not implemented")


ContentTypeType = Literal[
    "application/json",
    "multipart/form-data",
    "application/x-www-form-urlencoded",
    "application/yaml",
]


class RequestBody(BaseModel):
    model_config = ConfigDict(protected_namespaces=())

    description: Optional[str]

    """Exploded properties from the schema"""
    properties: List[Parameter] = []

    content_type: Optional[ContentTypeType]
    required: Optional[bool]

    model_name: Optional[str]


class Operation(BaseModel):
    tag: str
    name: str
    url: str
    method: str

    description: Optional[str] = None
    # path parameters
    path_parameters: Optional[List[Parameter]] = None
    query_parameters: Optional[List[Parameter]] = None

    # e.g. CreateUser
    request_body: Optional[RequestBody] = None
    response_model: Optional[Response] = None


def get(obj: Any, path: List[Union[str, int, None]], default: Any = None) -> Any:
    """Get by path.

    >>> assert get([{"a": {"b": 2}}], [0, "a", "b"]) == 2
    """
    # Took inspiration from funcy
    for k in path:
        try:
            obj = obj[k]  # queryable like a dict
        except (KeyError, IndexError, TypeError):
            try:
                obj = getattr(obj, str(k))  # queryable like a class
            except AttributeError:
                return default
    return obj


def getter(path: List[Union[str, int, None]]) -> Callable[[Any], Any]:
    def g(obj: Any) -> Any:
        return get(obj, path)

    return g


def get_schema_name(ref: Optional[str] = None) -> Optional[str]:
    # '#/components/schemas/JobResponse' -> JobResponse
    #
    # We use the ref to extract the name of the model
    if not ref:
        return None
    return ref.split("/")[-1]


def get_operations(paths: Any, model_schemas: dict[str, Schema]) -> List[Operation]:
    """Get all operations from all the path items."""
    temp_items: ItemsView[str, PathItem] = paths.items()

    operations = []
    for url, path in temp_items:
        methods = list(
            filter(lambda r: bool(r[1]), map(lambda m: (m, getattr(path, m)), METHODS))
        )

        operations.append(
            _get_operations_from_path(url, methods, model_schemas=model_schemas)
        )

    return list(chain.from_iterable(operations))


def _get_operations_from_path(
    url: str,
    methods: List[Tuple[str, OpenAPIOperation]],
    model_schemas: Dict[str, Schema],
) -> List[Operation]:
    """Get all operations from a single PathItem."""
    operations = []
    for method, op in methods:
        logger.debug(f"generating operation {op.operationId}", operation=op)
        path_parameters = extract_parameters(
            op.parameters or [], model_schemas, ParameterKind.PATH
        )
        query_parameters = extract_parameters(
            op.parameters or [], model_schemas, ParameterKind.QUERY
        )
        request_body = RequestBody(
            description=get(op, ["requestBody", "description"]),
            properties=get_properties(op, model_schemas),
            content_type=get_request_content_type(op),
            required=get(op, ["requestBody", "required"]),
            model_name=get_schema_name(_extract_request_schema_ref(op)),
        )
        operation = Operation(
            tag=op.tags[0],  # type: ignore[index]
            name=op.operationId,
            url=url,
            method=method,
            description=op.description,
            path_parameters=path_parameters,
            query_parameters=query_parameters,
            request_body=request_body,
            response_model=get_response_schema(op, model_schemas),
        )

        operations.append(operation)

    return operations


def get_schema_by_reference(reference: str, model_schemas: Dict[str, Schema]) -> Schema:
    """Get the schema from the reference.

    '#/components/schemas/JobResponse' -> Schema of JobResponse

    """
    schema_name = get_schema_name(reference)

    if not schema_name:
        raise SchemaNotFound(f"Schema name not found in reference {reference}")

    if schema_name not in model_schemas:
        raise SchemaNotFound(f"Schema {schema_name} not found in model schemas")

    return model_schemas[schema_name]


def get_properties(
    operation: OpenAPIOperation, model_schemas: Dict[str, Schema]
) -> List[Parameter]:
    content_type = get_request_content_type(operation)

    # Handle inline schemas (arrays, primitives) for application/json
    if content_type == "application/json":
        # Check if it's an inline schema (not a $ref)
        schema = get(
            operation,
            [
                "requestBody",
                "content",
                "application/json",
                "media_type_schema",
            ],
        )

        if schema and not get(schema, ["ref"]):
            # It's an inline schema, not a reference
            schema_type = get(schema, ["type"])
            title = get(schema, ["title"])

            if schema_type and title:
                param_name = title.lower()

                # Determine the Python type
                if schema_type.value == "array":
                    items_type = get(schema, ["items", "type"])
                    if items_type:
                        python_type = f"List[{to_python_type(items_type.value)}]"
                    else:
                        python_type = "List[Any]"
                else:
                    python_type = to_python_type(schema_type.value)

                param = Parameter(
                    name=param_name,
                    type=python_type,
                    required=True,
                    kind=ParameterKind.QUERY,  # Using QUERY as a generic kind
                    is_enum=False,
                    is_python_type=True,  # The type is already in Python format
                )
                return [param]

    if content_type != "multipart/form-data":
        # We currently only use properties when we need to add additional
        # form fields alongside a file upload.
        return []

    ref = _extract_request_schema_ref(operation)

    if not ref:
        return []

    try:
        schema = get_schema_by_reference(ref, model_schemas)
    except SchemaNotFound:
        return []

    if not schema.properties:
        return []

    parameters: List[Parameter] = []
    for property_name, property_metadata in schema.properties.items():
        if isinstance(property_metadata, openapi_pydantic.Reference):
            raise NotImplementedError(
                "We do not support nested references in request bodies"
            )
        is_required = (
            False if schema.required is None else property_name in schema.required
        )

        _type: str

        if isinstance(property_metadata.type, list):
            raise NotImplementedError(
                "We do not support multiple types in request bodies"
            )

        if property_metadata.type:
            _type = property_metadata.type.value
        else:
            if property_metadata.anyOf:
                # If it's not in `required`, we consider that it is Optional
                # with the only other type being the non-optional type.
                possible_types = [
                    t.type for t in property_metadata.anyOf if isinstance(t, Schema)
                ]

                unpacked = (
                    possible_types[0][0]
                    if isinstance(possible_types[0], list)
                    else possible_types[0]
                )

                if not unpacked:
                    raise ValueError("Expected type, got None.")

                _type = unpacked.value

            else:
                # FIXME: We do not support any other type than 'anyOf'
                # When 'allOf' is used, it's generally an Enum,
                # and for the uses in which this function is used,
                # we do not care about enums
                continue
        parameter = Parameter(
            required=is_required,
            name=property_name,
            type=_type,
            kind=ParameterKind.FORM,
        )
        parameters.append(parameter)

    return parameters


def extract_parameters(
    parameters: List[OpenApiParameter | Reference],
    model_schemas: Dict[str, Schema],
    kind: ParameterKind,
) -> List[Parameter]:
    def extract_parameter_of_kind(
        p: OpenApiParameter | Reference,
    ) -> TypeGuard[OpenApiParameter]:
        if isinstance(p, Reference):
            return False
        return p.param_in == kind.value

    filtered: list[OpenApiParameter] = list(
        filter(extract_parameter_of_kind, parameters or [])
    )

    result: List[Parameter] = []
    for parameter in filtered:
        logger.debug(f"extracting parameter {parameter.name}", parameter=parameter)
        if isinstance(parameter.param_schema, Schema):
            # if the 'type' field is not None, the type is not an union nor optional
            is_unique_type = parameter.param_schema.type is not None
            if is_unique_type:
                result.append(
                    Parameter(
                        name=parameter.name,
                        type=parameter.param_schema.type,
                        required=bool(parameter.required),
                        kind=kind,
                    )
                )
            else:
                # FIXME: The following code does not handle non-optional Unions,
                # for example Union[int, string],
                # it only handles one specific type and None,
                # for exampleOptional[int] / Union[None, int]

                # If there are any objects of type Reference we convert them to Schema
                only_schemas = map(
                    (
                        lambda t: (
                            get_schema_by_reference(t.ref, model_schemas)
                            if isinstance(t, Reference)
                            else t
                        )
                    ),
                    parameter.param_schema.anyOf,  # type: ignore[arg-type]
                )

                first_non_nullable_schema = next(
                    filter(lambda schema: schema.type is not None, only_schemas)
                )
                schema = first_non_nullable_schema
                type_ = schema.title if schema.enum else schema.type
                result.append(
                    Parameter(
                        name=parameter.name,
                        type=type_,
                        required=bool(schema.required),
                        is_enum=bool(schema.enum),
                        kind=kind,
                    )
                )

        elif isinstance(parameter.param_schema, Reference):
            schema = get_schema_by_reference(parameter.param_schema.ref, model_schemas)

            if not schema.enum:
                raise NotImplementedError(
                    "We do not support non-enum references in parameters"
                )

            result.append(
                Parameter(
                    name=parameter.name,
                    type=schema.title,
                    required=bool(schema.required),
                    is_enum=True,
                    kind=kind,
                )
            )
        else:
            raise TypeError(f"Unknown type {type(parameter.param_schema)}")

    return result


def get_response_schema(
    op: OpenAPIOperation, model_schemas: Dict[str, Schema]
) -> Optional[Response]:
    """Get the  schema from the (optional) response from the operation.

    "responses": {
        "200": {
        "description": "Successful Response",
        "content": {
            "application/json": {
            "schema": {
                "title": "Response Users-Find Users",
                "$ref": "#/components/schemas/User"
            }
            }
        }
        },
    },
    """
    type_: Enum = get(
        op,
        [
            "responses",
            "200",
            "content",
            "application/json",
            "media_type_schema",
            "type",
        ],
    )

    ref = _extract_response_schema_ref(op, type_.value if type_ else None)

    if not ref:
        return None

    return Response(
        ref=ref, name=get_schema_name(ref), type=type_.value if type_ else None
    )


def _extract_response_schema_ref(
    op: OpenAPIOperation, type: Literal["array"], status_code: int = 200
) -> Optional[str]:
    """Get the schema name from the response of the operation."""
    if type and type == "array":
        result = get(
            op,
            [
                "responses",
                str(status_code),
                "content",
                "application/json",
                "media_type_schema",
                "items",
                "ref",
            ],
        )

        return cast(Optional[str], result)
    result = get(
        op,
        [
            "responses",
            str(status_code),
            "content",
            "application/json",
            "media_type_schema",
            "ref",
        ],
    )

    return cast(Optional[str], result)


def _extract_request_schema_ref(operation: OpenAPIOperation) -> Optional[str]:
    """Get the full schema ref path from the request body of the operation.

    The request body can be optional as well as the ref. For both of these cases, we
    return None.

    Example
    -------
    "operationId": "create_job",
    "requestBody": {
      "content": {
        "application/json": {
          "schema": {
            "$ref": "#/components/schemas/JobCreate"  <---------------
          }
        }
      }
    """
    result = get(
        operation,
        [
            "requestBody",
            "content",
            get_request_content_type(operation),
            "media_type_schema",
            "ref",
        ],
    )
    return cast(Optional[str], result)


def get_request_content_type(operation: OpenAPIOperation) -> Optional[ContentTypeType]:
    """Get content type of the (optional) request body from the operation.

    "operationId": "create_job",
    "requestBody": {
      "content": {
        "application/json": {  <---------------------- "application/json"
          "schema": {
            "$ref": "#/components/schemas/JobCreate"
          }
        }
      }
    """
    if not get(operation, ["requestBody"]):
        return None

    content_type: Optional[ContentTypeType] = list(
        get(
            operation,
            [
                "requestBody",
                "content",
            ],
        ).keys()
    )[0]

    return content_type


def generate(input_path: Path, output_path: Path) -> str:
    """Generate the models using datamodel-codegen.

    Note that we are using the cli rather than the python model
    as the latter is buggy.
    """
    command = [
        "datamodel-codegen",
        "--input",
        str(input_path),
        "--output",
        str(output_path),
        "--input-file-type",
        "openapi",
        "--output-model-type",
        "pydantic_v2.BaseModel",
        "--use-schema-description",
        # "--collapse-root-models",
        # "--field-constraints",
        "--set-default-enum-member",
        "--use-annotated",
        "--disable-timestamp",
    ]
    subprocess.check_output(command)  # nosec

    return Path(output_path).read_text()


def main(
    schema: Path = typer.Argument(None),
    output_dir: Path = typer.Option(..., file_okay=False, dir_okay=True, writable=True),
    log_level: str = typer.Option("INFO"),
    typecheck: bool = typer.Option(True),
) -> None:
    """Generate API client."""
    structlog.configure(
        wrapper_class=structlog.make_filtering_bound_logger(
            logging.getLevelName(log_level.upper())
        ),
    )

    class Template(Enum):
        API = "api_autogenerated_new.py.jinja2"
        CLIENT = "client_new.py.jinja2"
        MODELS = "models_new.py.jinja2"

    TEMPLATE_TO_OUTPUT = {
        Template.API: "api.py",
        Template.CLIENT: "client.py",
        Template.MODELS: "models.py",
    }

    def get_output_path(template: Template) -> Path:
        return output_dir / TEMPLATE_TO_OUTPUT[template]

    with open(schema) as file:
        formatted_schema = json.load(file)

    openapi = OpenAPI.model_validate(formatted_schema)

    paths = openapi.paths
    model_schemas = cast(dict[str, Schema], openapi.components.schemas)  # type: ignore[union-attr]
    if not paths:
        raise ValueError("Couldn't extract paths from openapi schema")

    operations = get_operations(paths, model_schemas=model_schemas)

    models_to_import: List[str] = sorted(
        set(
            filter(
                None,
                chain(
                    map(getter(["response_model", "name"]), operations),
                    map(getter(["request_body", "model_name"]), operations),
                ),
            )
        )
    )

    api_autogenerated_rendered = template_env.get_template(Template.API.value).render(
        operations=operations,
        models=models_to_import,
        api_version=openapi.info.version,
    )

    client_rendered = template_env.get_template(Template.CLIENT.value).render(
        tags=sorted({op.tag for op in operations}),
        api_version=openapi.info.version,
    )

    with TemporaryDirectory() as temporary_directory_name:
        temporary_directory = Path(temporary_directory_name)
        input_file = Path(temporary_directory_name) / "input.json"
        input_file.write_text(json.dumps(formatted_schema, indent=2))
        output_path = Path(temporary_directory / TEMPLATE_TO_OUTPUT[Template.MODELS])

        generated_models = generate(
            input_path=input_file,
            output_path=output_path,
        )

    models_rendered = template_env.get_template(Template.MODELS.value).render(
        api_version=openapi.info.version,
        generated_code=generated_models,
    )

    for template, content in zip(
        [Template.API, Template.CLIENT, Template.MODELS],
        [api_autogenerated_rendered, client_rendered, models_rendered],
    ):
        with open(get_output_path(template), "w") as f:
            f.write(content)

    for template in [
        Template.API,
        Template.CLIENT,
        Template.MODELS,
    ]:
        # Format and fix imports with ruff
        subprocess.run(
            ["uvx", "ruff@latest", "check", "--fix", str(get_output_path(template))],
            check=True,
        )
        subprocess.run(
            ["uvx", "ruff@latest", "format", str(get_output_path(template))],
            check=True,
        )

    # Run typechecking from the `avatar-python` folder with it's custom configuration
    # We remove the VIRTUAL_ENV variable to force uv to use the env from output_dir
    # This way, we "remove" all remaining traces of the current uv/pipenv infos
    # which generates some stubs conflicts otherwise
    # (as the generator & avatar-py have different dependencies)
    if typecheck:
        new_env = os.environ.copy()
        new_env.pop("VIRTUAL_ENV", None)

        # Try justfile first, fallback to Makefile
        if (output_dir.parent / "justfile").exists():
            subprocess.run(
                ["just", "-d", str(output_dir.parent), "typecheck"],
                check=True,
                env=new_env,
            )
        elif (output_dir.parent / "Makefile").exists():
            subprocess.run(
                f"make -C {output_dir.parent} typecheck",
                shell=True,  # nosec B602
                check=True,
                env=new_env,
            )


if __name__ == "__main__":
    typer.run(main)
