import { OpenAPIRegistry, OpenApiGeneratorV31, extendZodWithOpenApi } from "@asteasolutions/zod-to-openapi";
import { z } from "zod";

import { ContractResponse } from "./contract.js";
import { createErrorSchema } from "./http-response.js";

import type { ContractGuard } from "./contract-guard.js";
import type {
    Contract,
    ContractRoute,
    CreateContractContext,
    CreateContractContextData,
    createContract,
} from "./contract.js";
import type { IRouteLocationInterface } from "./route.js";
import type { RouteConfig } from "@asteasolutions/zod-to-openapi";

type ExtendDocument = (document: ReturnType<OpenApiGeneratorV31["generateDocument"]>) => void;

type OpenApiContract = Contract<IRouteLocationInterface & ContractRoute>;

/**
 * Create an OpenAPI documentation in JSON format from all available contracts.
 *
 * You need to pass all routes as an object or as a direct array of contracts.
 *
 * It only considers contracts which are named like `createContract*`.
 *
 * It automatically parses the `versions` of your contracts and creates multiple OpenAPI schemas as needed.
 *
 * You can also pass a filter function to filter the contracts before they are added to the OpenAPI document. By
 * default, all contracts are included which do not have the `internal` profile.
 *
 * @example
 * ```ts
 * import * as routes from "./route/index.js";
 * createOpenApiDocuments(routes);
 * ```
 */
function createOpenApiDocuments(
    routes: Record<string, any> | OpenApiContract[],
    useContext?: CreateContractContextData,
    extendDocument?: ExtendDocument,
    filter?: (contract: OpenApiContract) => boolean,
): Record<string, ReturnType<OpenApiGeneratorV31["generateDocument"]>> {
    const versionToContracts = new Map<string, OpenApiContract[]>();
    let useContracts: OpenApiContract[] = [];

    if (Array.isArray(routes)) {
        useContracts = routes;
    } else {
        for (const value of Object.values(routes)) {
            if (
                typeof value === "function" &&
                value.name.startsWith("createContract") &&
                !value.name.startsWith("createContractGuard")
            ) {
                const createContractFunction = value as () => ReturnType<typeof createContract>;
                const contract = createContractFunction();

                if (!contract.versions) {
                    // This is not a contract
                    continue;
                }

                useContracts.push(contract);
            }
        }
    }

    for (const contract of useContracts) {
        if (filter && !filter(contract)) {
            continue;
        } else if (!filter && contract.route.profiles?.includes("internal")) {
            continue;
        }

        for (const version of contract.versions) {
            const contracts = versionToContracts.get(version) || [];
            contracts.push(contract);
            versionToContracts.set(version, contracts);
        }
    }

    const versionToDocument = new Map<string, ReturnType<OpenApiGeneratorV31["generateDocument"]>>();
    for (const [version, contracts] of versionToContracts) {
        const document = createOpenApiDocumentForVersion(version, contracts, useContext, extendDocument);
        versionToDocument.set(version, document);
    }

    return Object.fromEntries(versionToDocument);
}

function createOpenApiDocumentForVersion(
    version: string,
    contracts: OpenApiContract[],
    useContext?: CreateContractContextData,
    extendDocument?: ExtendDocument,
    automaticallyFindInfiniteRecursion = true,
) {
    const { registry, context, errorSchema } = createRegistry(useContext);
    const documentObj: ReturnType<OpenApiGeneratorV31["generateDocument"]> = {
        info: {
            title: "API",
            version: version || "0.0.1",
        },
        openapi: "3.1.0",
        tags: [],
    };

    for (const contract of contracts) {
        contract.storage.run(context, () => {
            registerContract(contract, registry, errorSchema, documentObj);
        });
    }

    if (documentObj["x-tagGroups"]?.length) {
        documentObj["x-tagGroups"].sort((a, b) => a.name.localeCompare(b.name));

        for (const tagGroup of documentObj["x-tagGroups"]) {
            tagGroup.tags.sort((a, b) => a.localeCompare(b));
        }
    }

    if (!documentObj["x-tagGroups"]?.length) {
        delete documentObj["x-tagGroups"];
    }

    try {
        const openApi = new OpenApiGeneratorV31(registry.definitions);
        const document = openApi.generateDocument(documentObj);
        fixDocumentMultipartFormData(document);
        extendDocument?.(document);
        return document;
    } catch (e) {
        // Automatically try to find out which contract is causing an infinite recursion
        if (
            process.env.NODE_ENV !== "production" &&
            automaticallyFindInfiniteRecursion &&
            e instanceof Error &&
            e.message.includes("Maximum call stack size exceeded")
        ) {
            const failedItems: string[] = [];
            for (const contract of contracts) {
                const contractIdentifier = `${contract.route.method} ${contract.route.path}`;
                try {
                    createOpenApiDocumentForVersion(version, [contract], useContext, extendDocument, false);
                } catch (e) {
                    // Recursion found here
                    if (e instanceof Error && e.message.includes("Maximum call stack size exceeded")) {
                        failedItems.push(contractIdentifier);
                    }
                }
            }

            if (failedItems.length > 0) {
                console.log(
                    "\x1b[38;5;208m[Development only] Infinite recursion detected while generating the OpenAPI document. Trying to find the contract causing the issue...\x1b[0m",
                );
                for (const item of failedItems) {
                    const isLastItem = item === failedItems[failedItems.length - 1];
                    const prefix = isLastItem ? "└──" : "├──";
                    const underlinedItem = `\x1b[4m${item}\x1b[0m`;
                    console.log(
                        `${prefix} \x1b[31m${underlinedItem} \x1b[31mis causing the infinite recursion. Please fix it.\x1b[0m`,
                    );
                }
            } else {
                throw e;
            }
        }

        throw e;
    }
}

/**
 * When using a `format: binary` for a component within a request body, automatically switch from
 * `application/json` to `multipart/form-data`.
 *
 * @see https://swagger.io/docs/specification/v3_0/describing-request-body/multipart-requests/
 * @see https://zod.dev/json-schema#file-schemas
 */
function fixDocumentMultipartFormData(document: ReturnType<OpenApiGeneratorV31["generateDocument"]>) {
    for (const path of Object.values(document.paths)) {
        for (const method of Object.values(path)) {
            if (
                method.requestBody?.content?.["application/json"] &&
                JSON.stringify(method.requestBody.content["application/json"].schema).includes(`"format":"binary"`)
            ) {
                method.requestBody.content["multipart/form-data"] = method.requestBody.content["application/json"];
                delete method.requestBody.content["application/json"];
            }
        }
    }
}

function createRegistry(useContext?: CreateContractContextData) {
    const registry = new OpenAPIRegistry();
    const context: CreateContractContextData = useContext || {
        t: ((translation) => `${translation}`) as CreateContractContextData["t"],
    };
    const errorSchema = createErrorSchema();

    return { registry, context, errorSchema };
}

function registerContract(
    contract: OpenApiContract,
    registry: OpenAPIRegistry,
    errorSchema: any,
    document: ReturnType<OpenApiGeneratorV31["generateDocument"]> & {
        "x-tagGroups"?: Array<{
            name: string;
            tags: string[];
        }>;
    },
) {
    const { z: contractZod, storage, route } = contract;
    const guards = contract.guards();

    if (!document["x-tagGroups"]) {
        document["x-tagGroups"] = [];
    }

    /**
     * @see https://guides.scalar.com/scalar/scalar-api-references/openapi#openapi-specification__x-taggroups
     */
    const xTagGroups = document["x-tagGroups"];

    // This is needed to make params / query parameters work with the meta function:
    // https://github.com/asteasolutions/zod-to-openapi/issues/324#issue-3244586153
    // https://github.com/asteasolutions/zod-to-openapi/issues/330
    [z, contractZod, ...Object.values(guards).map((guard) => guard.z)].forEach(extendZodWithOpenApi);

    const routeDetails = contract.routeDetails();
    const params = contract.params();
    const headers = contract.headers();

    const openApiResponses: RouteConfig["responses"] = {};

    const { security, removeFromHeaders, removeFromParams } = addSecuritySchemesFromGuards(
        storage,
        guards,
        openApiResponses,
        registry,
        contractZod,
    );
    addResponsesFromContract(contract, openApiResponses, errorSchema, contractZod);

    // Create the request body if some exists
    let requestBody: RouteConfig["request"]["body"];
    const requestBodySchema = contract.request();
    if (!(requestBodySchema instanceof contractZod.ZodVoid)) {
        requestBody = {
            required: !requestBodySchema.isOptional(), // see also https://github.com/colinhacks/zod/issues/4812#issuecomment-3047900832
            content: {
                "application/json": {
                    schema: requestBodySchema as z.ZodObject<any>,
                },
            },
        };
    }

    const pathTags: string[] = [];
    for (const group of routeDetails.groups ?? []) {
        pathTags.push(typeof group === "string" ? group : group.name);

        if (typeof group === "object") {
            // Check if tag already exists to avoid duplicates
            const existingTag = document.tags.find((tag) => tag.name === group.name);
            if (!existingTag) {
                document.tags.push({
                    name: group.name,
                    description: group.description,
                });
            }

            const { parent } = group;

            // When using `x-tagGroups` once, every tagged route needs to be inside a parent tag group
            // as otherwise they will not be shown in the documentation.
            const useParent = parent || "Uncategorized";
            const parentTagGroupObj = xTagGroups.find((tag) => tag.name === useParent);
            if (parentTagGroupObj) {
                parentTagGroupObj.tags.push(group.name);
            } else {
                xTagGroups.push({
                    name: useParent,
                    tags: [group.name],
                });
            }
        }
    }

    // Unique the tag groups to avoid duplicates in the sidebar
    for (const tagGroup of xTagGroups) {
        tagGroup.tags = [...new Set(tagGroup.tags)];
    }

    const useHeaders =
        headers instanceof contractZod.ZodObject && removeFromHeaders.length > 0
            ? headers.omit(
                  removeFromHeaders.reduce(
                      (acc, key) => {
                          acc[key] = true;
                          return acc;
                      },
                      {} as Record<string, true>,
                  ),
              )
            : headers;
    const useParams =
        params instanceof contractZod.ZodObject && removeFromParams.length > 0
            ? params.omit(
                  removeFromParams.reduce(
                      (acc, key) => {
                          acc[key] = true;
                          return acc;
                      },
                      {} as Record<string, true>,
                  ),
              )
            : params;

    registry[routeDetails.trigger === "webhook" ? "registerWebhook" : "registerPath"]({
        method: contract.route.method.toLowerCase() as RouteConfig["method"],
        path: contract.route.path,
        summary: routeDetails.shortDescription,
        tags: pathTags,
        "x-badges": route.profiles
            ? route.profiles.map((profile) => ({ name: profile, position: "before", color: "#ffcc00" }))
            : undefined,
        description: routeDetails.longDescription,
        security,
        request: {
            // @ts-expect-error this works as expected but zod v4 typing is a mismatch with the OpenAPI specification defined in `zod-to-openapi`.
            query: useParams,
            // @ts-expect-error this works as expected but zod v4 typing is a mismatch with the OpenAPI specification defined in `zod-to-openapi`.
            headers: useHeaders,
            body: requestBody,
        },
        responses: openApiResponses,
    });
}

function addResponsesFromContract(
    contract: OpenApiContract,
    openApiResponses: RouteConfig["responses"],
    errorSchema: any,
    zod: typeof z,
): void {
    const responses = contract.response();

    for (const [statusCode, response] of Object.entries({
        ...responses,
        default: ContractResponse.fromSchema({
            schema: zod.array(errorSchema),
        }),
    })) {
        openApiResponses[statusCode] = {
            description: response.description || "",
        };

        openApiResponses[statusCode].content = {
            "application/json": {
                schema: response.schema instanceof zod.ZodVoid ? zod.null() : response.schema,
            },
        };
    }
}

function addSecuritySchemesFromGuards(
    storage: CreateContractContext,
    guards: Record<string, ContractGuard>,
    openApiResponses: RouteConfig["responses"],
    registry: OpenAPIRegistry,
    zod: typeof z,
) {
    const security: RouteConfig["security"] = [];
    const removeFromHeaders = new Set<string>();
    const removeFromParams = new Set<string>();

    for (const guard of Object.values(guards)) {
        for (const type of ["params", "headers"] as const) {
            const guardSchemas = guard[type](storage);
            if (guardSchemas instanceof zod.ZodVoid) {
                continue;
            }

            for (const [key, value] of Object.entries(guardSchemas.shape)) {
                const meta = value.meta();
                if (meta?.isSecuritySchema) {
                    registry.registerComponent("securitySchemes", guard.name, {
                        type: "apiKey",
                        description: meta.description,
                        in: type === "params" ? "query" : "header",
                        name: key,
                    });
                    security.push({
                        [guard.name]: [],
                    });

                    // We need to remove the header or param from the request body as they are configured as security schemes
                    if (type === "headers") {
                        removeFromHeaders.add(key);
                    } else {
                        removeFromParams.add(key);
                    }
                }
            }
        }

        // Add the responses of the guard to the responses of the contract
        for (const [statusCode, response] of Object.entries(guard.response(storage))) {
            if (
                // When a guard is successful, those responses are never returned to the client
                +statusCode >= 200 &&
                +statusCode < 300
            ) {
                continue;
            }

            openApiResponses[statusCode] = {
                description: response.description || "",
            };

            openApiResponses[statusCode].content = {
                "application/json": {
                    schema: response.schema instanceof zod.ZodVoid ? zod.null() : response.schema,
                },
            };
        }
    }

    return {
        security,
        removeFromHeaders: Array.from(removeFromHeaders),
        removeFromParams: Array.from(removeFromParams),
    };
}

export { createOpenApiDocuments };
