from __future__ import annotations

import argparse
import subprocess
from datetime import datetime
from pathlib import Path


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Generate Pydantic schemas")
    parser.add_argument(
        "--spec-path",
        default="src/rabbitmq_mcp/data/rabbitmq-http-api-openapi.yaml",
    )
    parser.add_argument(
        "--output-path",
        default="src/rabbitmq_mcp/schemas/generated_schemas.py",
    )
    parser.add_argument("--force", action="store_true")
    return parser.parse_args()


def should_regenerate(spec_path: Path, output_path: Path, force: bool) -> bool:
    if force or not output_path.exists():
        return True
    return spec_path.stat().st_mtime > output_path.stat().st_mtime


def inject_rabbitmq_validators(output_path: Path) -> None:
    content = output_path.read_text(encoding="utf-8")

    if "import re" not in content:
        lines = content.split("\n")
        for i, line in enumerate(lines):
            if line.startswith("from pydantic import"):
                lines.insert(i, "import re")
                break
        content = "\n".join(lines)

    if "field_validator" not in content:
        content = content.replace(
            "from pydantic import", "from pydantic import field_validator,"
        )

    queue_validators = {
        "name": """
    @field_validator('name', mode='after')
    @classmethod
    def validate_queue_name(cls, v: str) -> str:
        if not v or len(v) > 255:
            raise ValueError("Queue name must be 1-255 characters")
        if not re.match(r'^[a-zA-Z0-9._-]+$', v):
            raise ValueError("Queue name must contain only alphanumeric, underscore, dash, dot")
        return v
""",
        "vhost": """
    @field_validator('vhost', mode='after')
    @classmethod
    def validate_vhost(cls, v: str | None) -> str:
        if not v:
            return "/"
        if not re.match(r'^[a-zA-Z0-9/_-]+$', v):
            raise ValueError("Vhost must contain only URL-safe characters")
        return v
""",
        "durable": """
    @field_validator('durable', mode='after')
    @classmethod
    def validate_durable(cls, v: bool | None) -> bool:
        if v is None:
            return True
        if not isinstance(v, bool):
            raise TypeError("Durable must be a boolean")
        return v
""",
    }

    exchange_validators = {
        "name": """
    @field_validator('name', mode='after')
    @classmethod
    def validate_exchange_name(cls, v: str) -> str:
        if v == "":
            return v
        if len(v) > 255:
            raise ValueError("Exchange name must be 0-255 characters")
        if not re.match(r'^[a-zA-Z0-9._-]+$', v):
            raise ValueError("Exchange name must contain only alphanumeric, underscore, dash, dot")
        return v
""",
        "type": """
    @field_validator('type', mode='after')
    @classmethod
    def validate_exchange_type(cls, v: str) -> str:
        valid_types = {"direct", "fanout", "topic", "headers"}
        if v not in valid_types:
            raise ValueError(f"Exchange type must be one of: {', '.join(valid_types)}")
        return v
""",
        "vhost": queue_validators["vhost"],
        "durable": queue_validators["durable"],
    }

    entity_validators = {
        "Queue": queue_validators,
        "Exchange": exchange_validators,
    }

    lines = content.split("\n")
    result_lines: list[str] = []
    i = 0

    while i < len(lines):
        line = lines[i]
        result_lines.append(line)
        if line.strip().startswith("class ") and "(BaseModel)" in line:
            class_name = line.strip().split()[1].rstrip(":").split("(")[0]
            indent_level = len(line) - len(line.lstrip())
            class_lines = [line]
            i += 1

            while i < len(lines):
                next_line = lines[i]
                if (
                    next_line.strip()
                    and not next_line.startswith(" " * (indent_level + 1))
                    and not next_line.strip().startswith("#")
                ):
                    break
                class_lines.append(next_line)
                i += 1

            class_text = "\n".join(class_lines)
            validators_to_add: list[str] = []
            applicable_validators = None
            for entity_pattern, validators_dict in entity_validators.items():
                if entity_pattern.lower() in class_name.lower():
                    applicable_validators = validators_dict
                    break

            if applicable_validators:
                for field_name, validator_code in applicable_validators.items():
                    field_pattern = f"{field_name}:"
                    if field_pattern in class_text or f"{field_name} =" in class_text:
                        validator_func_name = "validate_" + (
                            "exchange_type"
                            if field_name == "type" and "Exchange" in class_name
                            else (
                                "exchange_name"
                                if field_name == "name" and "Exchange" in class_name
                                else ("queue_name" if field_name == "name" else field_name)
                            )
                        )
                        if validator_func_name not in class_text:
                            validators_to_add.append(validator_code)

            if validators_to_add:
                result_lines.pop()
                result_lines.extend(class_lines[:-1])
                for validator in validators_to_add:
                    result_lines.append(validator)
                if class_lines[-1].strip() == "":
                    result_lines.append("")
            else:
                result_lines.extend(class_lines[1:])
            continue
        i += 1

    output_path.write_text("\n".join(result_lines), encoding="utf-8")


def main() -> None:
    args = parse_args()
    spec_path = Path(args.spec_path)
    output_path = Path(args.output_path)
    if not should_regenerate(spec_path, output_path, args.force):
        print("Generated schemas are up-to-date")
        return

    output_path.parent.mkdir(parents=True, exist_ok=True)

    result = subprocess.run(
        [
            "datamodel-codegen",
            "--input",
            str(spec_path),
            "--output",
            str(output_path),
            "--input-file-type",
            "openapi",
            "--output-model-type",
            "pydantic_v2.BaseModel",
            "--use-standard-collections",
            "--use-schema-description",
            "--field-constraints",
            "--use-default",
            "--strict-nullable",
            "--use-annotated",
            "--snake-case-field",
            "--enable-faux-immutability",
        ],
        capture_output=True,
        text=True,
    )
    if result.returncode != 0:
        raise SystemExit(result.stderr)

    header = (
        "# Auto-generated from OpenAPI - DO NOT EDIT MANUALLY\n"
        f"# Generated: {datetime.now().isoformat()}\n"
        f"# Source: {spec_path}\n\n"
    )
    output_path.write_text(header + output_path.read_text(encoding="utf-8"), encoding="utf-8")
    inject_rabbitmq_validators(output_path)
    print(f"Generated schemas at {output_path}")


if __name__ == "__main__":
    main()
