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

import click
from jinja2 import Environment, FileSystemLoader
from linkml_runtime.linkml_model.meta import Element, SlotDefinition
from linkml_runtime.utils.formatutils import camelcase, underscore
from linkml_runtime.utils.schemaview import SchemaView

from linkml.generators.docgen import DocGenerator
from linkml.utils.generator import Generator, shared_arguments


@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: Optional[str] = None  # output directory with generated markdown files
    template_file: Optional[str] = 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()
