import argparse
import keyword
import os
import pathlib
import re
from typing import Any

from django.core.management.base import BaseCommand

from esi import __build_date__, __title__, __url__, __version__, stubs
from esi.openapi_clients import ESIClientProvider


def sanitize_class_name(name: str) -> str:
    """Convert a tag into a valid Python class name."""
    sanitized = re.sub(r'[^0-9a-zA-Z_]', '_', name.strip())
    if sanitized and not sanitized[0].isalpha():
        sanitized = f"_{sanitized}"
    return sanitized


def sanitize_operation_class(name: str) -> str:
    return re.sub(r'[^0-9a-zA-Z_]', '', name[0].upper() + name[1:] + "Operation")


class ModelGenerator:
    """Generate Pydantic model stubs from OpenAPI schemas and return type names.

    This only affects typing in the generated stubs (pyi) and does not change runtime behavior.
    """

    def __init__(self, components: dict | None = None) -> None:
        self._models: dict[str, list[str]] = {}
        self._name_counts: dict[str, int] = {}
        self._schema_cache: dict[int, str] = {}
        self._components = components or {}
        self._component_class_map: dict[str, str] = {}
        self._aliases: dict[str, str] = {}

    def _unique_class_name(self, base: str) -> str:
        base = sanitize_class_name(base)
        if base not in self._models and base not in self._name_counts:
            self._name_counts[base] = 1
            return base
        # Ensure uniqueness by suffixing
        n = self._name_counts.get(base, 1)
        while True:
            candidate = f"{base}_{n}"
            if candidate not in self._models:
                self._name_counts[base] = n + 1
                return candidate
            n += 1

    def _singleton_model(self, name: str, lines: list[str]) -> None:
        """Only emit once per class name"""
        if name in self._models:
            return
        self._models[name] = lines

    def _singleton_alias(self, name: str, rhs: str) -> None:
        """Only emit once per alias"""
        if name in self._models:
            return
        # Avoid recursive aliases like `UUID = UUID`
        if name == rhs:
            return
        self._models[name] = [f"{name} = {rhs}"]
        self._aliases[name] = rhs

    def _properties_dict(self, schema) -> dict:
        """.properties or ._properties"""
        props = getattr(schema, "properties", None)
        if props is None:
            props = getattr(schema, "_properties", None)
        return props or {}

    def _required_list(self, schema) -> list[str]:
        """.required or ._required"""
        req = getattr(schema, "required", None)
        if req is None:
            req = getattr(schema, "_required", None)
        if isinstance(req, (list, tuple, set)):
            return list(req)
        return []

    def _additional_properties(self, schema) -> Any | None:
        return getattr(schema, "additionalProperties", None)

    def _enum_values(self, schema) -> list[str] | None:
        """.enum or ._enum"""
        values = getattr(schema, "enum", None)
        if values is None:
            values = getattr(schema, "_enum", None)
        return list(values) if isinstance(values, (list, tuple, set)) else None

    def _format(self, schema) -> str | None:
        """.format or ._format"""
        fmt = getattr(schema, "format", None)
        if fmt is None:
            fmt = getattr(schema, "_format", None)
        return fmt

    def _ref_string(self, schema) -> str | None:
        ref = getattr(schema, "ref", None)
        if isinstance(ref, str):
            return ref
        root = getattr(schema, "root", None)
        if isinstance(root, dict) and "$ref" in root:
            return root.get("$ref")
        return None

    def _get_component(self, name: str) -> Any | None:
        comps = self._components or {}
        # components may be an object with .schemas as dict-like
        if hasattr(comps, "get"):
            return comps.get(name)
        schemas = getattr(comps, "schemas", None)
        if isinstance(schemas, dict):
            return schemas.get(name)
        return None

    def _mapping_alias_str(self, additional_schema: Any, name_hint: str) -> str:
        """Build a dict alias type string for additionalProperties.

        Returns a string like 'dict[str, Any]' or 'dict[str, Schema]'
        """
        if additional_schema is True or additional_schema is None:
            return "dict[str, Any]"
        if additional_schema is False:
            # No mapping allowed; caller decides how to handle
            return ""
        inner = self.schema_to_type(additional_schema, f"{name_hint}Value")
        return f"dict[str, {inner}]"

    def _base_primitive_type(self, schema_type: str | None, fmt: str | None) -> str:
        """Mapping of OpenAPI primitive+format to Python base type."""
        if schema_type == "string":
            match fmt:
                case "date-time":
                    return "datetime"
                case "date":
                    return "date"
                case "uuid":
                    return "UUID"
                case _:
                    return "str"
        if schema_type == "integer":
            return "int"
        if schema_type == "number":
            return "float"
        if schema_type == "boolean":
            return "bool"
        if schema_type in {"array", "object"}:
            # handled elsewhere
            return "Any"
        return "Any"

    def schema_to_type(self, schema: Any, name_hint: str = "Model") -> str:
        """Convert an OpenAPI schema to a Python type hint, generating Pydantic models for objects."""
        if not schema:
            return "Any"

        # Cache repeated schemas (by identity) to avoid duplicate class generation
        cache_key = id(schema)
        if cache_key in self._schema_cache:
            return self._schema_cache[cache_key]

        # Try to resolve common cases early
        schema_type = getattr(schema, "type", None)

        # Handle nullable across styles (nullable flag OR type includes null OR anyOf/oneOf with null)
        nullable_flag = bool(getattr(schema, "nullable", False))
        if isinstance(schema_type, list):
            if "null" in schema_type:
                # Remove null and keep underlying type for processing
                schema_type = next((t for t in schema_type if t != "null"), None)
                nullable_flag = True

        def _wrap_nullable(type_str: str) -> str:
            if nullable_flag and type_str != "None" and not type_str.endswith(" | None"):
                return f"{type_str} | None"
            return type_str

        # Resolve $ref components
        ref = self._ref_string(schema)
        if isinstance(ref, str) and ref.startswith("#/components/schemas/"):
            comp_name = ref.split("/")[-1]
            # If we already mapped this component, reuse
            if comp_name in self._component_class_map:
                result = self._component_class_map[comp_name]
                self._schema_cache[cache_key] = result
                return result
            comp_schema = self._get_component(comp_name)
            # If the component schema is available, build a proper type
            if comp_schema is not None:
                # Prefer using component name as class name
                class_name = sanitize_class_name(comp_name)
                enum_vals = self._enum_values(comp_schema)
                comp_type = getattr(comp_schema, "type", None)
                if enum_vals and comp_type in {"string", "integer", "number"}:
                    lines = [f"class {class_name}(Enum):"]
                    # Normalize enum member names
                    for idx, val in enumerate(enum_vals):
                        if isinstance(val, str):
                            member = re.sub(r"\W+", "_", val).upper()
                            if not member or member[0].isdigit():
                                member = f"VALUE_{idx}"
                            lines.append(f"    {member} = '{val}'")
                        else:
                            lines.append(f"    VALUE_{idx} = {val}")
                    self._singleton_model(class_name, lines)
                    self._component_class_map[comp_name] = class_name
                    self._schema_cache[cache_key] = class_name
                    return class_name

                # For arrays and primitives, emit type aliases instead of models
                props = self._properties_dict(comp_schema)
                if comp_type == "array":
                    inner = self.schema_to_type(getattr(comp_schema, "items", None), f"{class_name}Item")
                    alias_rhs = f"list[{inner}]"
                    alias_rhs = alias_rhs if not getattr(comp_schema, "nullable", False) else f"{alias_rhs} | None"
                    self._singleton_alias(class_name, alias_rhs)
                    self._component_class_map[comp_name] = class_name
                    self._schema_cache[cache_key] = class_name
                    return class_name
                if comp_type == "object" and not props:
                    # If object with no explicit properties, use a mapping alias for flexibility
                    ap = self._additional_properties(comp_schema)
                    alias_rhs = self._mapping_alias_str(ap, class_name) if ap is not False else class_name
                    if alias_rhs and alias_rhs != class_name:
                        if getattr(comp_schema, "nullable", False):
                            alias_rhs = f"{alias_rhs} | None"
                        self._singleton_alias(class_name, alias_rhs)
                        self._component_class_map[comp_name] = class_name
                        self._schema_cache[cache_key] = class_name
                        return class_name
                if comp_type in {"integer", "number", "string", "boolean"} and not props:
                    # Respect known string formats via centralized mapper
                    fmt = self._format(comp_schema)
                    base = self._base_primitive_type(comp_type, fmt)
                    alias_rhs = base if not getattr(comp_schema, "nullable", False) else f"{base} | None"
                    self._singleton_alias(class_name, alias_rhs)
                    self._component_class_map[comp_name] = class_name
                    self._schema_cache[cache_key] = class_name
                    return class_name

                # Otherwise, defer to normal handling as an object and force class name
                t = self._object_schema_to_model(comp_schema, class_name)
                self._component_class_map[comp_name] = t
                self._schema_cache[cache_key] = t
                return t

        # Handle composed schemas
        one_of = getattr(schema, "oneOf", None)
        any_of = getattr(schema, "anyOf", None)
        all_of = getattr(schema, "allOf", None)
        if one_of or any_of:
            raw_variants = one_of or any_of
            variants = raw_variants if isinstance(raw_variants, (list, tuple)) else []
            if not variants:
                return "Any"
            types: list[str] = []
            local_nullable = nullable_flag
            for idx, subschema in enumerate(variants):
                sub_type = getattr(subschema, "type", None)
                if sub_type == "null" or (isinstance(sub_type, list) and "null" in sub_type):
                    local_nullable = True
                    continue
                t = self.schema_to_type(subschema, f"{name_hint}Alt{idx}")
                if t not in types:
                    types.append(t)
            if not types:
                return "Any"
            union = " | ".join(types)
            if local_nullable and "None" not in union:
                union = f"{union} | None"
            # Add discriminator hint comment when available
            disc = getattr(schema, "discriminator", None)
            disc_name = None
            if disc is not None:
                disc_name = getattr(disc, "propertyName", None) or getattr(disc, "property_name", None)
                if disc_name is None:
                    disc_root = getattr(disc, "root", None)
                    if isinstance(disc_root, dict):
                        disc_name = disc_root.get("propertyName")
            if disc_name:
                which = "oneOf" if one_of else "anyOf"
                union = f"{union}  # {which} discriminator: {disc_name}"
            self._schema_cache[cache_key] = union
            return union
        if all_of:
            # Attempt to merge object schemas in allOf
            class_name = self._unique_class_name(f"{name_hint}")
            props: dict = {}
            required: set[str] = set()
            additional_schema = None
            for subschema in all_of:
                sub_props = self._properties_dict(subschema)
                if sub_props:
                    for k, v in sub_props.items():
                        props[k] = v
                sub_required = self._required_list(subschema)
                required.update(sub_required)
                ap = self._additional_properties(subschema)
                if ap is not None:
                    additional_schema = ap
            if not props and additional_schema is None:
                # Fall back
                return "Any"
            # Build a synthetic schema-like object with combined properties

            class Dummy:
                pass
            dummy = Dummy()
            setattr(dummy, "properties", props)
            setattr(dummy, "required", list(required))
            if additional_schema is not None:
                setattr(dummy, "additionalProperties", additional_schema)
            result = self._object_schema_to_model(dummy, class_name)
            result = _wrap_nullable(result)
            self._schema_cache[cache_key] = result
            return result

        if schema_type == "array":
            items_schema = getattr(schema, "items", None)
            inner = self.schema_to_type(items_schema, f"{name_hint}Item")
            return _wrap_nullable(f"list[{inner}]")

        if schema_type == "object" or self._properties_dict(schema):
            props = self._properties_dict(schema)
            additional = self._additional_properties(schema)
            # If no explicit properties but has additionalProperties, treat as mapping
            if not props and additional is not None:
                alias_rhs = self._mapping_alias_str(additional, name_hint)
                if alias_rhs:
                    return _wrap_nullable(alias_rhs)

            class_name = self._unique_class_name(f"{name_hint}")
            result = self._object_schema_to_model(schema, class_name)
            result = _wrap_nullable(result)
            self._schema_cache[cache_key] = result
            return result

        if isinstance(schema_type, str):
            enum_vals = self._enum_values(schema)
            if enum_vals:
                lit_vals: list[str] = []
                for v in enum_vals:
                    if isinstance(v, str):
                        lit_vals.append(repr(v))
                    else:
                        lit_vals.append(str(v))
                return _wrap_nullable(f"Literal[{', '.join(lit_vals)}]")
            fmt = self._format(schema)
            base = self._base_primitive_type(schema_type, fmt)

            # Add constrained annotations for common string constraints
            if schema_type == "string":
                min_len = getattr(schema, "minLength", None) or getattr(schema, "_minLength", None)
                max_len = getattr(schema, "maxLength", None) or getattr(schema, "_maxLength", None)
                pattern = getattr(schema, "pattern", None) or getattr(schema, "_pattern", None)
                if pattern is None:
                    root = getattr(schema, "root", None)
                    if isinstance(root, dict):
                        min_len = min_len or root.get("minLength")
                        max_len = max_len or root.get("maxLength")
                        pattern = root.get("pattern") or pattern
                field_args: list[str] = []
                if isinstance(min_len, int):
                    field_args.append(f"min_length={min_len}")
                if isinstance(max_len, int):
                    field_args.append(f"max_length={max_len}")
                if isinstance(pattern, str) and pattern:
                    field_args.append(f"pattern={repr(pattern)}")
                if field_args:
                    annotated = f"Annotated[{base}, Field({', '.join(field_args)})]"
                    return _wrap_nullable(annotated)
            return _wrap_nullable(base)

        return _wrap_nullable("Any")  # Fallback

    def write_models(self, f) -> None:
        """Write all collected models to file-like f in declaration order."""
        if not self._models:
            return
        for name, lines in self._models.items():
            for line in lines:
                f.write(line + "\n")
            f.write("\n\n")

    def _object_schema_to_model(self, schema: Any, class_name: str) -> str:
        # If class already emitted, reuse
        if class_name in self._models:
            return class_name

        lines: list[str] = [f"class {class_name}(BaseModel):"]

        props = self._properties_dict(schema)
        if not props:
            lines.append("    pass")
        else:
            required = set(self._required_list(schema))
            for prop_name, prop_schema in props.items():
                original = str(prop_name)
                py_name = re.sub(r"[^0-9a-zA-Z_]", "_", original)
                if not py_name:
                    py_name = "_field"
                # Ensure starts with letter or underscore
                if not (py_name[0].isalpha() or py_name[0] == "_"):
                    py_name = f"_{py_name}"
                # Avoid keywords and invalid identifiers
                if keyword.iskeyword(py_name) or not py_name.isidentifier():
                    py_name = f"{py_name}_"
                sub_hint = f"{class_name}_{py_name[0].upper() + py_name[1:]}"
                typ = self.schema_to_type(prop_schema, sub_hint)
                if original not in required:
                    typ = f"{typ} | None"
                # Add alias when sanitized name differs
                if py_name != original:
                    typ_repr = f"Annotated[{typ}, Field(alias='{original}')]"
                else:
                    typ_repr = typ
                lines.append(f"    {py_name}: {typ_repr}")

        self._singleton_model(class_name, lines)
        return class_name


class Command(BaseCommand):
    help = "Generate ESI Stubs from the current ESI spec with correct type hints."

    def add_arguments(self, parser: argparse.ArgumentParser) -> None:
        parser.add_argument(
            "--output",
            default=None,
            help="Custom output path for the generated ESI stub file (default: stubs.pyi next to openapi_clients.py).",
        )
        parser.add_argument(
            "--compatibility_date",
            default=__build_date__,
            help="Compatibility Date to build ESI Stubs from.",
        )
        parser.add_argument(
            "--spec_file",
            required=False,
            type=pathlib.Path,
            default=None,
            help="Path to a OpenApi 3.1 spec file to load."
        )

    def handle(self, *args: Any, **options: Any) -> None:
        self.stdout.write("Starting ESI stub generation...")

        if spec_file := options["spec_file"]:
            cwd = os.getcwd()
            spec_file = os.path.join(cwd, spec_file)

        stub_api = ESIClientProvider(
            ua_appname=__title__,
            ua_url=__url__,
            ua_version=__version__,
            compatibility_date=options["compatibility_date"],
            spec_file=spec_file,
        )

        stub_api = stub_api.client.api

        base_path = os.path.dirname(stubs.__file__)
        output_path = options["output"] or os.path.join(base_path, "stubs.pyi")

        self.stdout.write(f"Outputting to: {output_path}")

        spec_root = stub_api._root
        with open(output_path, "w", encoding="utf-8") as f:
            # Local helpers to reduce duplication
            def _get_json_body_schema(op_obj: Any) -> Any | None:
                rb = getattr(op_obj, "requestBody", None)
                if not rb:
                    return None
                rb_content = getattr(rb, "content", {})
                if isinstance(rb_content, dict):
                    json_media = rb_content.get("application/json")
                    if json_media is not None:
                        return getattr(json_media, "schema_", None)
                return None

            def _clean_doc(op_obj: Any) -> str:
                text = getattr(op_obj, "description", None) or getattr(op_obj, "summary", None) or ""
                return text.replace("\n", " ").strip()

            def _get_response_json_schema(op_obj: Any, method: str) -> Any | None:
                """Return the JSON response schema object for an operation and HTTP method.

                post -> 201, others -> 200, put/delete return None.
                """
                try:
                    if method == "post":
                        resp = op_obj.responses.get("201")
                    elif method in ("put", "delete"):
                        return None
                    else:
                        resp = op_obj.responses.get("200")
                    content = getattr(resp, "content", {}) if resp else {}
                    if isinstance(content, dict) and "application/json" in content:
                        return content["application/json"].schema_
                except Exception:
                    return None
                return None

            def _is_paginated(op_obj: Any) -> bool:
                """Check whether an operation supports page-based pagination."""
                return any(getattr(p, "name", None) in ("before", "after", "page") for p in getattr(op_obj, "parameters", []))

            # File headers
            f.write("# flake8: noqa=E501\n")
            f.write("# cSpell: disable\n")
            f.write("# Auto Generated do not edit\n")
            # Python Imports
            f.write("from datetime import date, datetime\n")
            f.write("from enum import Enum\n")
            f.write("from typing import Annotated, Any, Literal\n")
            f.write("from uuid import UUID\n\n")
            f.write("from pydantic import BaseModel, Field\n\n")
            f.write("from esi.models import Token\n")
            f.write("from esi.openapi_clients import EsiOperation\n\n\n")

            operation_classes = {}
            # Attempt to get OpenAPI components for $ref resolution
            components = getattr(getattr(stub_api, "_root", None), "components", None)
            schemas = getattr(components, "schemas", None) if components else None
            model_gen = ModelGenerator(components=schemas if isinstance(schemas, dict) else None)

            # Create the current common error Schema in ESI currently
            if isinstance(schemas, dict) and "Error" in schemas:
                model_gen.schema_to_type(schemas["Error"], "Error")

            # Pre-collect request body models/aliases so they are written out before client methods reference them
            for tag in sorted(stub_api._operationindex._tags.keys()):
                ops = stub_api._operationindex._tags[tag]
                for nm, op in sorted(ops._operations.items()):
                    op_obj = op[2]
                    schema_obj = _get_json_body_schema(op_obj)
                    if schema_obj is not None:
                        body_name = f"{sanitize_operation_class(nm)}Body"
                        # Ensure a named model/alias is generated for request bodies
                        t = model_gen.schema_to_type(schema_obj, body_name)
                        # If schema maps to a non-object primitive/array and no class was created with body_name,
                        # make a type alias so the name exists in the stubs.
                        if t != body_name and body_name not in model_gen._models:
                            model_gen._singleton_alias(body_name, t)

            for tag in sorted(stub_api._operationindex._tags.keys()):
                # ESI Operation
                # The various methods called on an ESI Operation
                # result(), Results(), Results_Localized() etc. all live here
                ops = stub_api._operationindex._tags[tag]
                for nm, op in sorted(ops._operations.items()):
                    op_type = op[0]
                    op_obj = op[2]
                    docstring = _clean_doc(op_obj)
                    op_class_name = sanitize_operation_class(nm)

                    response_type = "Any"
                    schema = _get_response_json_schema(op_obj, op_type)
                    if op_type in ("put", "delete"):
                        response_type = "None"
                    elif schema is not None:
                        response_type = model_gen.schema_to_type(schema, f"{op_class_name}Response")

                    # If the response is an alias to a list type (e.g. Foo = list[Bar]) we want
                    # results() to return the flattened list[Bar], not list[list[Bar]]
                    if response_type.startswith("list["):
                        inner_list_type = response_type
                        # response_type already a list[...] so results() should also be list[...] (same type)
                        results_type = inner_list_type
                    else:
                        # Maybe it's an alias to list[...] recorded earlier
                        alias_rhs = getattr(model_gen, "_aliases", {}).get(response_type)
                        if isinstance(alias_rhs, str) and alias_rhs.startswith("list["):
                            results_type = alias_rhs
                        else:
                            results_type = f"list[{response_type}]"

                    if op_class_name not in operation_classes:
                        f.write(f"class {op_class_name}(EsiOperation):\n")
                        if response_type != "None" and _is_paginated(op_obj):
                            f.write("    \"\"\"EsiOperation, use result(), results() or results_localized()\"\"\"\n")
                        else:
                            f.write("    \"\"\"EsiOperation, use result()\"\"\"\n")

                        # result()
                        f.write(f"    def result(self, use_etag: bool = True, return_response: bool = False, force_refresh: bool = False, use_cache: bool = True, store_cache: bool = True, **extra) -> {response_type}:\n")  # noqa: E501
                        f.write(f"        \"\"\"{docstring}\"\"\"\n") if docstring else None
                        f.write("        ...\n\n")
                        if response_type != "None" and _is_paginated(op_obj):
                            # We only need the extra utility functions if its actually an endpoint that returns data
                            # results()
                            f.write(f"    def results(self, use_etag: bool = True, return_response: bool = False, force_refresh: bool = False, use_cache: bool = True, store_cache: bool = True, **extra) -> {results_type}:\n")  # noqa: E501
                            f.write(f"        \"\"\"{docstring}\"\"\"\n") if docstring else None
                            f.write("        ...\n\n")

                            # results_localized()
                            f.write(f"    def results_localized(self, languages: list[str] | str | None = None, **extra) -> dict[str, {results_type}]:\n")  # noqa: E501
                            f.write(f"        \"\"\"{docstring}\"\"\"\n") if docstring else None
                            f.write("        ...\n\n\n")

                        operation_classes[op_class_name] = True

            # Emit collected Pydantic models before operations/classes so types are available
            model_gen.write_models(f)

            f.write("class ESIClientStub:\n")
            tags = sorted(stub_api._operationindex._tags.keys())
            for idx, tag in enumerate(tags):
                # ESI ESIClientStub
                # The various ESI Tags and Operations
                ops = stub_api._operationindex._tags[tag]
                class_name = f"_{sanitize_class_name(tag)}"
                f.write(f"    class {class_name}:\n")
                for nm, op in sorted(ops._operations.items()):
                    op_obj = op[2]
                    op_class_name = sanitize_operation_class(nm)
                    effective_security = (
                        getattr(op_obj, "security", None) or getattr(spec_root, "security", None)
                    )

                    def _has_oauth2(sec: Any) -> bool:
                        data = sec if isinstance(sec, dict) else getattr(sec, "root", None)
                        if not isinstance(data, dict):
                            return False
                        return any(k.lower() == "oauth2" for k in data)

                    needs_oauth = any(_has_oauth2(s) for s in (effective_security or []))

                    params = ["self"]
                    optional_params = []
                    schema_obj = _get_json_body_schema(op_obj)
                    if schema_obj is not None:
                        type_hint = model_gen.schema_to_type(schema_obj, f"{op_class_name}Body")
                        params.append(f"body: {type_hint}")

                    for p in getattr(op_obj, "parameters", []):
                        required = getattr(p, "required", False)
                        schema_obj = getattr(p, "schema_", None)
                        if schema_obj is not None:
                            # Use generator for richer types (arrays/objects)
                            param_type = model_gen.schema_to_type(schema_obj, f"{op_class_name}_{p.name}")
                        else:
                            param_type = "Any"
                        default = ""
                        if not required:
                            param_type = f"{param_type} | None"
                            default = " = ..."
                        param_name = p.name.replace("-", "_")
                        if param_name == "authorization" and needs_oauth:
                            # Skip the Authorization Parameter, we inject this at HTTP Header Level
                            continue
                        if required:
                            params.append(f"{param_name}: {param_type}{default}")
                        else:
                            optional_params.append(f"{param_name}: {param_type}{default}")
                    if needs_oauth:
                        # Here, we add our own custom param instead of the earlier Authorization
                        # Our library will pick this up and use it later
                        params.append("token: Token")
                    optional_params.append("**kwargs: Any")
                    params_str = ", ".join(params + optional_params)
                    docstring = _clean_doc(op_obj)

                    if docstring:
                        f.write(f"        def {nm}({params_str}) -> {op_class_name}:\n")
                        f.write(f"            \"\"\"{docstring}\"\"\"\n")
                        f.write("            ...\n\n")
                    else:
                        f.write(f"        def {nm}({params_str}) -> {op_class_name}: ...\n")

                # Only add a single trailing newline after the final tag as this is the end of the file as well then.
                end_newlines = "\n" if idx == len(tags) - 1 else "\n\n"
                f.write(f"\n    {sanitize_class_name(tag)}: {class_name} = {class_name}()" + end_newlines)

        self.stdout.write(self.style.SUCCESS(f"ESI stubs written to {output_path}"))
