import importlib
import logging
import os
from dataclasses import dataclass, field
from pathlib import Path

import click
from jinja2 import Environment, FileSystemLoader

from linkml.generators.docgen import DocGenerator
from linkml.utils.generator import Generator, shared_arguments
from linkml_runtime.linkml_model.meta import Element, SlotDefinition
from linkml_runtime.utils.formatutils import camelcase, underscore
from linkml_runtime.utils.schemaview import SchemaView


@dataclass
class MermaidClassDiagramGenerator(Generator):
    """This generator creates Mermaid Class diagrams (https://mermaid.js.org/syntax/classDiagram.html)
    for individual classes in a given LinkML schema and outputs them as Markdown files. It uses
    a jinja template to render these class diagrams, which can be customized by the user and pointed
    to using the `--template-file` option. If no template is provided, then the default template
    file provided by LinkML and present at linkml/generators/docgen/class_diagram.md.jinja2 will be used.
    The generator also has the option of specifying a certain set of classes for which you want the
    class diagrams to be generated by using the `--classes` option.
    """

    generatorname = os.path.basename(__file__)
    generatorversion = "0.0.1"
    valid_formats = ["markdown"]
    uses_schemaloader = False
    requires_metamodel = False

    directory: str | None = None  # output directory with generated markdown files
    template_file: str | None = None  # custom/default jinja template for class diagrams
    classes: list[str] = field(default_factory=list)  # optional subset of classes
    preserve_names: bool = False  # preserve original LinkML names in diagram output

    def __post_init__(self):
        super().__post_init__()
        self.logger = logging.getLogger(__name__)
        self.schemaview = SchemaView(self.schema, merge_imports=self.mergeimports)

        # set the output directory
        self.output_directory = Path(self.directory)

        # set the template file
        if not self.template_file:
            package_dir = os.path.dirname(importlib.util.find_spec("linkml").origin)
            # Default location of the template file
            self.template_file = os.path.join(package_dir, "generators", "docgen", "class_diagram.md.jinja2")

    def generate_class_diagrams(self):
        """Generate Mermaid class diagrams for the specified subset of classes
        or all classes if none are specified.
        """
        self.output_directory.mkdir(parents=True, exist_ok=True)

        template_folder = os.path.dirname(self.template_file)
        template_name = os.path.basename(self.template_file)
        loader = FileSystemLoader(template_folder)
        env = Environment(loader=loader)
        temp_doc_gen = DocGenerator(self.schema, mergeimports=self.mergeimports, preserve_names=self.preserve_names)
        temp_doc_gen.customize_environment(env)

        template = env.get_template(template_name)

        all_classes = self.schemaview.all_classes()

        if self.classes:
            class_items = [(cn, all_classes[cn]) for cn in self.classes if cn in all_classes]
        else:
            class_items = list(all_classes.items())

        for cn, class_def in class_items:
            self.logger.info(f"Generating Mermaid diagram for class: {cn}")
            rendered = template.render(gen=self, element=class_def, schemaview=self.schemaview)
            filename = self.name(class_def) if self.preserve_names else cn
            outfile = self.output_directory / f"{filename}.md"
            with open(outfile, "w", encoding="utf-8") as f:
                f.write(rendered)

    def cardinality(self, slot: SlotDefinition) -> str:
        """Reuses the cardinality logic from DocGenerator."""
        return DocGenerator.cardinality(slot)

    def mermaid_directive(self) -> str:
        """Provides the code fence directive (e.g., `mermaid` or `{mermaid}`)."""
        return "mermaid"

    def name(self, element: Element) -> str:
        """Returns the canonical name for an element."""
        if element is None:
            return ""
        if self.preserve_names:
            return element.name
        elif type(element).class_name == "slot_definition":
            return underscore(element.name)
        elif type(element).class_name == "class_definition":
            return camelcase(element.name)
        else:
            return camelcase(element.name)

    def link_mermaid(self, element):
        """Generates a link for the given element."""
        return f"../{self.name(element)}"

    def all_type_object_names(self):
        return list(self.schemaview.all_types().keys())


@shared_arguments(MermaidClassDiagramGenerator)
@click.command()
@click.option(
    "--template-file",
    "-t",
    type=click.Path(exists=True),
    default=None,
    help="Path to Jinja template for class diagrams.",
)
@click.option(
    "--directory",
    "-d",
    type=click.Path(),
    required=True,
    help="Folder in which to write the Markdown files.",
)
@click.option(
    "--classes",
    "-c",
    multiple=True,
    help="One or more classes in the schema for which to generate diagrams. "
    "If omitted, diagrams for all classes are generated.",
)
@click.option(
    "--preserve-names/--normalize-names",
    default=False,
    show_default=True,
    help="Preserve original LinkML names in Mermaid diagram output (e.g., for class names, slot names, file names).",
)
@click.version_option(click.__version__, "-V", "--version")
def cli(yamlfile, template_file, directory, classes, **args):
    logging.basicConfig(level=logging.INFO)
    gen = MermaidClassDiagramGenerator(
        schema=yamlfile,
        template_file=template_file,
        directory=directory,
        classes=list(classes),  # convert tuple to a list
        **args,
    )
    gen.generate_class_diagrams()


if __name__ == "__main__":
    cli()
