"""
VisualizationGenerator - Document generator.
"""

from typing import TYPE_CHECKING, Optional

if TYPE_CHECKING:
    from ..analyzer import ProjectAnalyzer

from ..diagram_generator import DiagramConfig
from ..diagram_generator import DiagramGenerator as MermaidDiagramGenerator
from .base import BaseGenerator


class VisualizationGenerator(BaseGenerator):
    """Generator for visualization diagrams."""

    def __init__(self, analyzer: "ProjectAnalyzer", config: Optional[DiagramConfig] = None):
        """
        Initialize the visualization generator.

        Args:
            analyzer: ProjectAnalyzer instance
            config: Optional diagram configuration
        """
        super().__init__(analyzer)
        self.config = config or DiagramConfig()
        self.diagram_gen = MermaidDiagramGenerator(analyzer, self.config)

    def generate(self, diagram_type: Optional[str] = None) -> str:
        """
        Generate visualization diagrams.

        Args:
            diagram_type: Optional specific diagram type to generate.
                         If None, generates all diagrams.
                         Options: 'flowchart', 'sequence', 'component', 'er', 'architecture', 'all'

        Returns:
            Generated diagrams document
        """
        if diagram_type and diagram_type != "all":
            # Generate specific diagram type
            return self._generate_with_header(diagram_type)
        else:
            # Generate all diagrams
            return self.diagram_gen.generate_all_diagrams()

    def _generate_with_header(self, diagram_type: str) -> str:
        """Generate a specific diagram with header."""
        diagram_content = self.diagram_gen.generate_specific_diagram(diagram_type)

        header = f"# {diagram_type.capitalize()} Diagrams\n\n"
        header += f"**Generated**: {self.datetime}\n\n"
        header += "---\n\n"

        footer = "\n\n---\n\n*Generated by RE-cue - Business Process Visualization*\n"

        return header + diagram_content + footer
