#!/usr/bin/env python3
import argparse
from inspect import get_annotations, getmodule, getmro, isclass
from typing import get_type_hints

from blueapi.core.bluesky_types import is_bluesky_plan_generator
from blueapi.core.context import load_module_all
from pydantic import BaseModel

from mx_bluesky.common.parameters.components import MxBlueskyParameters
from mx_bluesky.hyperion import experiment_plans


def main():
    parser = argparse.ArgumentParser(
        description="Generate PlantUML diagram for Hyperion parameter models."
    )
    parser.add_argument(
        "outfile",
        help="Path to the output .puml file",
    )
    args = parser.parse_args()

    with open(args.outfile, "w") as out:
        write_header(out)

        parameter_types = [
            get_type_hints(obj)["parameters"]
            for obj in load_module_all(experiment_plans)
            if is_bluesky_plan_generator(obj)
        ]

        experiment_types = set(parameter_types)
        all_types = set(parameter_types)

        # Collect all BaseModel ancestors
        for experiment_type in experiment_types:
            for base in get_mro_base_models(experiment_type):
                all_types.add(base)

        # Group by package
        types_by_package = group_types_by_package(all_types)

        # Split mx_bluesky.common.parameters into mixins and main parameter types
        mx_bluesky_param_types, mixin_types = split_mx_bluesky_types(types_by_package)

        # Write mx_bluesky.common.parameters
        out.write("package mx_bluesky.common.parameters {\n")
        out.write("together {\n")
        for t in mixin_types:
            generate_class(t, out)
        out.write("}\n")
        out.write("together {\n")
        for t in mx_bluesky_param_types:
            generate_class(t, out)
        out.write("}\n")
        out.write("}\n")

        # Write mx_bluesky.hyperion.parameters
        out.write("package mx_bluesky.hyperion.parameters {\n")
        for t in types_by_package["mx_bluesky.hyperion.parameters"]:
            generate_class(t, out)
        out.write("}\n")

        # Write inheritance arrows
        for t in all_types:
            for base in t.__bases__:
                if base not in (object, BaseModel):
                    out.write(f"{base.__name__} <|-- {t.__name__}\n")

        out.write("@enduml\n")


def write_header(out):
    out.write(
        "'This file is auto-generated by generate_plantuml.py\n"
        "@startuml hyperion_parameter_model\n"
        "title Hyperion Parameter Model\n"
        "set namespaceSeparator none\n\n"
    )


def get_mro_base_models(t):
    for base in getmro(t):
        if issubclass(base, BaseModel) and base is not BaseModel:
            yield base


def group_types_by_package(all_types):
    grouped = {}
    for t in all_types:
        mod = getmodule(t)
        assert mod
        grouped.setdefault(mod.__package__, []).append(t)
    return grouped


def split_mx_bluesky_types(types_by_package):
    mx_bluesky_param_types = set()
    mixin_types = set()

    for t in types_by_package.get("mx_bluesky.common.parameters", []):
        if issubclass(t, MxBlueskyParameters):
            mx_bluesky_param_types.add(t)
        else:
            mixin_types.add(t)

    return mx_bluesky_param_types, mixin_types


def generate_class(t, out):
    out.write(f"class {t.__name__}{{\n")
    for field_name, field_type in get_annotations(t).items():
        out.write(f"\t{generate_type(field_type)} {field_name}\n")
    out.write("}\n")


def generate_type(field_type):
    return (
        field_type.__name__
        if isclass(field_type) and issubclass(field_type, BaseModel)
        else str(field_type)
    )


if __name__ == "__main__":
    main()
