# pyright: basic
import os
import re

from datamodel_code_generator.format import CustomCodeFormatter
from typing import Any

CONFIG_BODY_IDENT = "# config_body_start"
CAPITALIZE_WORDS = {
    "mp3",
    "pdf",
    "json",
    "lb",
    "url",
    "uri",
    "id",
    "ip",
    "db",
    "obs",
    "sms",
    "cn",
    "en",
    "zh",
    "sso",
    "html",
    "ivs",
}


def camel_to_splitted(name: str) -> list[str]:
    return (
        "".join([f"_{char.lower()}" if char.isupper() else char for char in name])
        .lstrip("_")
        .split("_")
    )


def snake_to_camel(name: str) -> str:
    splitted = [word for word in name.split("_")]
    for i, word in enumerate(splitted):
        if len(word) <= 2 or word.lower() in CAPITALIZE_WORDS:
            splitted[i] = word.upper()
        else:
            splitted[i] = word.capitalize()

    return "".join(splitted)


pydantic_import_re = re.compile(
    r"from pydantic import ([\w, ]+?)?(?:, )?BaseModel(?:, )?([\w, ]+)?"
)


class CodeFormatter(CustomCodeFormatter):
    config_file_ext: str
    config_name: str
    config_path: str
    config_type: str

    def __init__(self, formatter_kwargs: dict[str, Any]) -> None:
        super().__init__(formatter_kwargs)
        self.config_file_ext = formatter_kwargs["config_file_ext"]
        self.config_name = formatter_kwargs["config_name"]
        self.config_path = formatter_kwargs["config_path"]
        self.config_type = formatter_kwargs["config_type"]

    def apply(self, code: str) -> str:
        config_class_name = snake_to_camel(self.config_name)

        pydantic_import_match = pydantic_import_re.search(code)
        assert (
            pydantic_import_match is not None
        ), "datamodel-code-generator created unexpected result"

        pydantic_import_left, pydantic_import_right = pydantic_import_match.groups()

        code = pydantic_import_re.sub("from pydantic import BaseModel", code)

        if pydantic_import_left and pydantic_import_right:
            pydantic_import = (
                f"from pydantic import {pydantic_import_left}, {pydantic_import_right}"
            )
        elif pydantic_import_left:
            pydantic_import = f"from pydantic import {pydantic_import_left}"
        elif pydantic_import_right:
            pydantic_import = f"from pydantic import {pydantic_import_right}"
        else:
            pydantic_import = None

        new_imports = [
            f"from runch import RunchModel, RunchConfigReader",
            CONFIG_BODY_IDENT,
        ]

        if pydantic_import is not None:
            new_imports.insert(0, pydantic_import)

        code = code.replace(
            "from pydantic import BaseModel",
            "\n".join(new_imports),
        )
        code = code.replace("class Model(", f"class {config_class_name}ConfigModel(")
        code = code.replace("(BaseModel):", "(RunchModel):")

        config_head_start = code.find(
            f"class {config_class_name}ConfigModel(RunchModel):"
        )
        config_head = code[config_head_start:].strip("\n")

        config_body_start = code.find(CONFIG_BODY_IDENT)
        config_body = code[
            config_body_start + len(CONFIG_BODY_IDENT) : config_head_start
        ].strip("\n")

        config_imports = code[:config_body_start].strip("\n")

        config_renames: dict[str, str] = {}

        def transform_body_line(line: str) -> str:
            if line.startswith("class "):
                class_name = line.split("(")[0].split("class ")[1]
                class_name_splitted = camel_to_splitted(class_name)
                for i, word in enumerate(class_name_splitted):
                    if len(word) <= 2 or word in CAPITALIZE_WORDS:
                        class_name_splitted[i] = word.upper()
                    else:
                        class_name_splitted[i] = word.capitalize()

                # handle schema deduplication suffix generated by datamodel-codegen
                schema_dedup_id: str | None = None
                if re.match(r"^[A-Za-z]+\d+$", class_name_splitted[-1]):
                    _schema_dedup_id: str = re.findall(
                        r"\d+$", class_name_splitted[-1]
                    )[0]
                    schema_dedup_id = _schema_dedup_id
                    class_name_splitted[-1] = class_name_splitted[-1].replace(
                        _schema_dedup_id, f""
                    )

                new_class_name = "".join(class_name_splitted)
                config_renames[class_name] = new_class_name + "Config"
                if schema_dedup_id:
                    config_renames[class_name] += schema_dedup_id

            if line.strip() == "":
                return ""
            else:
                return line

        # add indent
        config_body = "\n".join(map(transform_body_line, config_body.split("\n")))
        config_classes = f"{config_head}\n\n{config_body}"

        for old_name, new_name in config_renames.items():
            old_name_pattern_1 = re.compile(rf" {old_name}\b")
            old_name_pattern_2 = re.compile(rf"\[{old_name}\b")
            config_classes = old_name_pattern_1.sub(f" {new_name}", config_classes)
            config_classes = old_name_pattern_2.sub(f"[{new_name}", config_classes)

        config_file_fullname = os.path.basename(self.config_path)
        config_file_path = os.path.dirname(self.config_path)

        config_read_directive = f'_{self.config_name}_reader = RunchConfigReader[{config_class_name}ConfigModel]("{config_file_fullname}", config_dir="{config_file_path}", config_type="{self.config_type}")'
        config_read_directive += (
            f"\n{self.config_name} = _{self.config_name}_reader.read_lazy()"
        )
        config_read_directive += (
            "\n\n# uncomment the following line to enable the watch_update feature"
        )
        config_read_directive += (
            f"\n# _{self.config_name}_reader"
            + '.set_feature("watch_update", {"enabled": True, "args": {"update_interval": 10}})'
        )

        return f"{config_imports}\n\n{config_classes}\n\n{config_read_directive}"
