import base64
import logging
import subprocess
from enum import Enum
from functools import reduce
from pathlib import Path
from typing import Self

from pydantic import BaseModel

from ._latex import model_to_table
from ._visitor import (
    Relation,
    all_classes_in_model,
    classes_in_model,
    related_classes_in_model,
    walk_model_to_depth,
)

logger = logging.getLogger(__name__)


__all__ = [
    "model_to_plantuml_class",
    "generate_plantuml_diagrams",
    "PlantUMLDiagram",
]

RELATION_MAP = {
    Relation.contains: "1",
    Relation.contains_many: "..*",
    Relation.contains_one_of: "0..1",
}

START_UML = "@startuml"
END_UML = "@enduml"


def model_name(model: type[BaseModel | Enum]) -> str:
    """Return name of model used for PlantUML.

    This will also include the optional namespace if it exists.
    """
    if hasattr(model, "_namespace"):
        return f"{model._namespace}.{model.__name__}"
    return model.__name__


def model_to_plantuml_class(model: type[BaseModel | Enum]) -> str:
    """Return PlantUML class string."""
    puml = [START_UML]
    if issubclass(model, Enum):
        puml.append(f"class {model_name(model)} <<enum>> {{")
    else:
        puml.append(f"class {model_name(model)} {{")

    if issubclass(model, BaseModel) or isinstance(model, BaseModel):
        table = model_to_table(model, optional_symbol=" [opt]", type_sep=" | ")
        for row in table:
            puml.append(f"    + {row['Name']}: <i>{row['Type']}</i>")
    elif issubclass(model, Enum):
        for member in model:
            puml.append(f"    + {member}")
    else:
        msg = f"couldn't convert: {model}"
        raise TypeError(msg)

    puml.append("}")
    puml.append(END_UML)

    return "\n".join(puml)


def _uml_relations_for_models(
    models: list[type[BaseModel]],
    allowed_types=BaseModel,
) -> tuple[list[str], set[type[BaseModel | Enum]]]:
    """
    Generate PlantUML inheritance relationships.

    The search goes both up (parents) and down (children), recursively.

    Parameters
    ----------
    models : list[type]
        Classes to analyze.
    allowed_types : type | tuple[type], optional
        Only include classes that are subclasses of these types.

    Returns
    -------
    list[str]
        Sorted list of PlantUML relationship lines like "Parent <|-- Child".
    set[type[BaseModel|Enum]]
        List if models that should be imported
    """
    # Start from the given classes and expand outward
    to_visit = set(models)
    all_related = set(models)
    seen = set()
    models_to_import = set()

    while to_visit:
        cls = to_visit.pop()
        if cls in seen:
            continue
        seen.add(cls)

        parents, children = related_classes_in_model(cls)

        new_related = set(parents) | set(children)
        to_visit |= new_related - all_related
        all_related |= new_related

    # Now build PlantUML relationships among all related classes
    relations = set()
    for cls in all_related:
        models_to_import.add(cls)
        for parent in cls.__mro__[1:]:
            if issubclass(parent, allowed_types) and parent in all_related:
                models_to_import.add(parent)
                relations.add(f"{model_name(parent)} <|-- {model_name(cls)}")

    return sorted(relations), models_to_import


def _generate_imports(models: list[type[BaseModel]]) -> list[str]:
    """Generate import statement for the models."""
    models = list(set(models))  # demove dups
    puml = []

    for model in models:
        puml.append(f"!include ./definitions/{model_name(model)}.plantuml")

    return puml


def model_to_plantuml_relations(
    model: type[BaseModel],
    *,
    max_depth=1000,
    import_definitions: bool = False,
    show_inheritance: bool = False,
) -> str:
    """
    Output PlantUML relationships for a model.

    By default containment relations are generated, and optionally inheritance
    relations.

    Parameters
    ----------
    model: type[BaseModel]
       the model
    max_depth: int
       how deep to go
    import_defintions: bool
       include plantuml imports of the class definitions
    show_inheritance: bool
       include parent/child relations

    Returns
    -------
    str:
        PlantUML code
    """
    puml = [START_UML]

    models_to_import = set()
    models_to_import.add(model)

    for _, path in walk_model_to_depth(model, max_depth=max_depth):
        if len(path) < 2:
            continue

        item = path[-1]
        parent = path[-2]

        if isinstance(item.item_type, type) and issubclass(
            item.item_type, BaseModel | Enum
        ):
            rel = RELATION_MAP[item.parent_relation]
            if item.is_optional and item.parent_relation != Relation.contains_one_of:
                rel = "0.." + rel

            models_to_import.add(item.item_type)

            puml.append(
                f'{model_name(parent.item_type)} "+{item.item_name}" *-- "{rel}"'
                f" {model_name(item.item_type)} "
            )

    if show_inheritance:
        relations, relation_models = _uml_relations_for_models([model])
        models_to_import.update(relation_models)
        puml += relations

    if import_definitions:
        puml += _generate_imports(list(models_to_import))

    puml.append(END_UML)

    return "\n".join(puml)


def generate_plantuml_diagrams(
    models: list[type[BaseModel]],
    output_dir: Path | str,
    show_inheritance: bool = True,
    imports: bool = True,
    max_depth: int = 1000,
) -> None:
    """
    Generate plantuml diagrams for the given models.

    Parameters
    ----------
    models: list[type[BaseModel]]
       list of Models to generte diagrams for
    output_dir: Path | str
       where to write the plantuml code
    show_inheritance: bool
       If True, also include parent/child relations and classes
    max_depth: int
       maximum depth to include in the diagrams.
    """
    base = Path(output_dir)
    defs = base / "definitions"
    defs.mkdir(parents=True, exist_ok=True)

    # First generate all necessary class definitions

    all_classes = []
    for model in models:
        all_classes += all_classes_in_model(model)

    all_classes = list(set(all_classes))

    for model in all_classes:
        def_filename = defs / f"{model_name(model)}.plantuml"
        def_filename.write_text(model_to_plantuml_class(model))
        logger.debug("Wrote: %s", def_filename)

    # now generate the relationship diagrams:
    logger.debug("Writing relationship diagrams...")
    for model in models:
        if imports:
            rel_filename = base / f"relations_{model_name(model)}.plantuml"
            rel_filename.write_text(
                model_to_plantuml_relations(
                    model, max_depth=max_depth, show_inheritance=show_inheritance
                )
            )
            logger.debug("Wrote: %s", rel_filename)
        full_filename = base / f"full_{model_name(model)}.plantuml"
        full_filename.write_text(
            model_to_plantuml_relations(
                model,
                import_definitions=True,
                max_depth=max_depth,
                show_inheritance=True,
            )
        )
        logger.debug("Wrote: %s", full_filename)


class PlantUMLDiagram:
    """Render a PlantUML diagram in a jupyter notebook.

    Diagrams can be composed using the ``+`` operator.
    """

    def __init__(
        self,
        model_or_text: str | type[BaseModel | Enum],
        relations: bool = True,
        inheritance: bool = False,
        details: bool = False,
        max_depth: int = 1000,
        cwd: Path | str | None = None,
    ) -> None:
        """
        Create a PlantUML diagram.

        Parameters
        ----------
        model_or_text: str | type[BaseModel]
            Either a model, or any arbitrary plantuml-format text.
        relations: bool
            include composed classes in the diagram
        inheritance: bool
            include inherited classes in the diagram
        details: bool
            include details of inherited or related classes
        cwd: Path | str
            current working directory, in case you have a relative
            ``!include`` directive in the plantuml text.
        """
        if isinstance(model_or_text, str):
            self.plantuml_text = model_or_text
        elif isinstance(model_or_text, type) and issubclass(
            model_or_text, BaseModel | Enum
        ):
            self.plantuml_text = model_to_plantuml_class(model_or_text)
            if relations and issubclass(model_or_text, BaseModel):
                self.plantuml_text = (
                    self
                    + PlantUMLDiagram(
                        model_to_plantuml_relations(
                            model_or_text,
                            show_inheritance=inheritance,
                            max_depth=max_depth,
                        )
                    )
                ).plantuml_text
            if details and issubclass(model_or_text, BaseModel):
                classes = classes_in_model(model_or_text, max_depth=max_depth)
                if len(classes) > 1:
                    self.plantuml_text = (
                        self
                        + reduce(
                            lambda x, y: x + y,
                            (
                                PlantUMLDiagram(m, inheritance=False, relations=False)
                                for m in classes
                                if m is not model_or_text
                            ),
                        )
                    ).plantuml_text
        else:
            raise ValueError("Expected plantuml string or type[BaseModel]")
        self.cwd = cwd

    @classmethod
    def from_path(cls, path) -> Self:
        """Return diagram from an existing file at path."""
        path = Path(path)
        return cls(path.read_text(), cwd=path.parent)

    def _repr_html_(self) -> str:
        try:
            png_data = self._generate_png_pipe()
            b64_data = base64.b64encode(png_data).decode("utf-8")
            return f'<img src="data:image/png;base64,{b64_data}" />'
        except Exception as e:
            return (
                '<div style="color: red;">Error rendering PlantUML diagram:'
                f" {str(e)}</div>"
            )

    def _generate_png_pipe(self) -> bytes:
        """Generate PNG using PlantUML's pipe mode."""
        result = subprocess.run(
            ["plantuml", "-tpng", "-pipe"],
            cwd=self.cwd,
            input=self.plantuml_text.encode("utf-8"),
            capture_output=True,
            check=True,
        )
        return result.stdout

    def __add__(self, other) -> Self:
        text = self.plantuml_text.replace(START_UML, "").replace(END_UML, "")
        other_text = other.plantuml_text.replace(START_UML, "").replace(END_UML, "")
        return PlantUMLDiagram(f"@startuml\n{text}\n{other_text}\n@enduml")

    def __str__(self) -> str:
        """Turn diagram to string."""
        return self.plantuml_text
