"""Export chunks to graph formats (GraphML, DOT)."""

from __future__ import annotations

import xml.etree.ElementTree as ET
from pathlib import Path
from typing import TYPE_CHECKING, Any
from xml.dom import minidom

from chunker.interfaces.export import (
    ChunkRelationship,
    ExportFormat,
    ExportMetadata,
    GraphExporter,
    RelationshipType,
)

if TYPE_CHECKING:
    import io
    from collections.abc import Iterator

    from chunker.types import CodeChunk


class GraphMLExporter(GraphExporter):
    """Export chunks and relationships to GraphML fmt."""

    def __init__(self):
        self._node_attributes: list[str] = [
            "node_type",
            "language",
            "file_path",
            "start_line",
            "end_line",
            "parent_context",
        ]
        self._edge_attributes: list[str] = ["relationship_type"]
        self._layout_algorithm: str | None = None

    def export(
        self,
        chunks: list[CodeChunk],
        relationships: list[ChunkRelationship],
        output: Path | io.IOBase,
        metadata: ExportMetadata | None = None,
    ) -> None:
        """Export chunks with relationships to GraphML.

        Args:
            chunks: List of code chunks
            relationships: List of chunk relationships
            output: Output path or stream
            metadata: Export metadata
        """
        graphml = self._build_graphml(chunks, relationships, metadata)
        xml_str = self._prettify_xml(ET.tostring(graphml, encoding="unicode"))
        if isinstance(output, str | Path):
            Path(output).write_text(xml_str, encoding="utf-8")
        else:
            output.write(xml_str)

    def export_streaming(
        self,
        chunk_iterator: Iterator[CodeChunk],
        relationship_iterator: Iterator[ChunkRelationship],
        output: Path | io.IOBase,
    ) -> None:
        """Export using iterators for large datasets.

        Note: GraphML requires the full graph structure, so we collect
        all data before writing. For true streaming, consider DOT fmt.
        """
        chunks = list(chunk_iterator)
        relationships = list(relationship_iterator)
        self.export(chunks, relationships, output)

    def set_node_attributes(self, attributes: list[str]) -> None:
        """Set which chunk attributes to include as node properties."""
        self._node_attributes = attributes

    def set_edge_attributes(self, attributes: list[str]) -> None:
        """Set which relationship attributes to include as edge properties."""
        self._edge_attributes = attributes

    def add_layout_hints(self, layout_algorithm: str) -> None:
        """Add layout hints for visualization."""
        self._layout_algorithm = layout_algorithm

    @staticmethod
    def supports_format(fmt: ExportFormat) -> bool:
        """Check if this exporter supports a fmt."""
        return fmt == ExportFormat.GRAPHML

    def get_schema(self) -> dict[str, Any]:
        """Get the export schema."""
        return {
            "fmt": "graphml",
            "version": "1.0",
            "node_attributes": self._node_attributes,
            "edge_attributes": self._edge_attributes,
            "layout_algorithm": self._layout_algorithm,
        }

    def _build_graphml(
        self,
        chunks: list[CodeChunk],
        relationships: list[ChunkRelationship],
        metadata: ExportMetadata | None,
    ) -> ET.Element:
        """Build GraphML XML structure."""
        graphml = ET.Element(
            "graphml",
            {
                "xmlns": "http://graphml.graphdrawing.org/xmlns",
                "xmlns:xsi": "http://www.w3.org/2001/XMLSchema-instance",
                "xsi:schemaLocation": "http://graphml.graphdrawing.org/xmlns http://graphml.graphdrawing.org/xmlns/1.0/graphml.xsd",
            },
        )
        if metadata:
            comment = f"Generated by TreeSitter Chunker - {metadata.created_at}"
            graphml.append(ET.Comment(comment))
        for i, attr in enumerate(self._node_attributes):
            key = ET.SubElement(
                graphml,
                "key",
                {
                    "id": f"d{i}",
                    "for": "node",
                    "attr.name": attr,
                    "attr.type": "string",
                },
            )
        edge_attr_offset = len(self._node_attributes)
        for i, attr in enumerate(self._edge_attributes):
            key = ET.SubElement(
                graphml,
                "key",
                {
                    "id": f"d{edge_attr_offset + i}",
                    "for": "edge",
                    "attr.name": attr,
                    "attr.type": "string",
                },
            )
        graph = ET.SubElement(graphml, "graph", {"id": "G", "edgedefault": "directed"})
        for chunk in chunks:
            node = ET.SubElement(graph, "node", {"id": chunk.chunk_id})
            for i, attr in enumerate(self._node_attributes):
                value = getattr(chunk, attr, "")
                if value:
                    data = ET.SubElement(node, "data", {"key": f"d{i}"})
                    data.text = str(value)
            if self._layout_algorithm:
                data = ET.SubElement(node, "data", {"key": "layout"})
                data.text = self._layout_algorithm
        for idx, rel in enumerate(relationships):
            edge = ET.SubElement(
                graph,
                "edge",
                {
                    "id": f"e{idx}",
                    "source": rel.source_chunk_id,
                    "target": rel.target_chunk_id,
                },
            )
            data = ET.SubElement(edge, "data", {"key": f"d{edge_attr_offset}"})
            data.text = rel.relationship_type.value
            if rel.metadata:
                for key, value in rel.metadata.items():
                    if key in self._edge_attributes[1:]:
                        attr_idx = self._edge_attributes.index(key)
                        data = ET.SubElement(
                            edge,
                            "data",
                            {"key": f"d{edge_attr_offset + attr_idx}"},
                        )
                        data.text = str(value)
        return graphml

    @staticmethod
    def _prettify_xml(xml_str: str) -> str:
        """Pretty print XML with proper indentation."""
        dom = minidom.parseString(xml_str)
        return dom.toprettyxml(indent="  ")


class DOTExporter(GraphExporter):
    """Export chunks and relationships to Graphviz DOT fmt."""

    def __init__(self):
        self._node_attributes: list[str] = ["node_type", "language"]
        self._edge_attributes: list[str] = ["relationship_type"]
        self._layout_algorithm: str = "dot"
        self._node_style: dict[str, str] = {
            "shape": "box",
            "style": "rounded,filled",
            "fillcolor": "lightblue",
        }
        self._edge_style: dict[str, dict[str, str]] = {
            RelationshipType.CALLS.value: {"color": "blue", "style": "solid"},
            RelationshipType.INHERITS.value: {
                "color": "red",
                "style": "solid",
                "arrowhead": "empty",
            },
            RelationshipType.IMPLEMENTS.value: {
                "color": "green",
                "style": "dashed",
                "arrowhead": "empty",
            },
            RelationshipType.IMPORTS.value: {"color": "gray", "style": "dotted"},
            RelationshipType.USES.value: {"color": "purple", "style": "dashed"},
            RelationshipType.DEPENDS_ON.value: {"color": "orange", "style": "dotted"},
        }

    def export(
        self,
        chunks: list[CodeChunk],
        relationships: list[ChunkRelationship],
        output: Path | io.IOBase,
        metadata: ExportMetadata | None = None,
    ) -> None:
        """Export chunks with relationships to DOT fmt.

        Args:
            chunks: List of code chunks
            relationships: List of chunk relationships
            output: Output path or stream
            metadata: Export metadata
        """
        dot_content = self._build_dot(chunks, relationships, metadata)
        if isinstance(output, str | Path):
            Path(output).write_text(dot_content, encoding="utf-8")
        else:
            output.write(dot_content)

    def export_streaming(
        self,
        chunk_iterator: Iterator[CodeChunk],
        relationship_iterator: Iterator[ChunkRelationship],
        output: Path | io.IOBase,
    ) -> None:
        """Export using iterators for large datasets."""
        if isinstance(output, str | Path):
            with Path(output).open("w", encoding="utf-8") as f:
                self._stream_dot(chunk_iterator, relationship_iterator, f)
        else:
            self._stream_dot(chunk_iterator, relationship_iterator, output)

    def set_node_attributes(self, attributes: list[str]) -> None:
        """Set which chunk attributes to include as node properties."""
        self._node_attributes = attributes

    def set_edge_attributes(self, attributes: list[str]) -> None:
        """Set which relationship attributes to include as edge properties."""
        self._edge_attributes = attributes

    def add_layout_hints(self, layout_algorithm: str) -> None:
        """Add layout hints for visualization."""
        self._layout_algorithm = layout_algorithm

    @staticmethod
    def supports_format(fmt: ExportFormat) -> bool:
        """Check if this exporter supports a fmt."""
        return fmt == ExportFormat.DOT

    def get_schema(self) -> dict[str, Any]:
        """Get the export schema."""
        return {
            "fmt": "dot",
            "version": "1.0",
            "node_attributes": self._node_attributes,
            "edge_attributes": self._edge_attributes,
            "layout_algorithm": self._layout_algorithm,
            "node_style": self._node_style,
            "edge_style": self._edge_style,
        }

    def set_node_style(self, style: dict[str, str]) -> None:
        """Set default node style attributes."""
        self._node_style.update(style)

    def set_edge_style(
        self,
        relationship_type: str,
        style: dict[str, str],
    ) -> None:
        """Set edge style for specific relationship type."""
        self._edge_style[relationship_type] = style

    def _build_dot(
        self,
        chunks: list[CodeChunk],
        relationships: list[ChunkRelationship],
        metadata: ExportMetadata | None,
    ) -> str:
        """Build DOT fmt content."""
        lines = []
        lines.append("digraph CodeStructure {")
        lines.append(f"  rankdir={self._layout_algorithm.upper()};")
        lines.append(
            "  node ["
            + ", ".join(f'{k}="{v}"' for k, v in self._node_style.items())
            + "];",
        )
        if metadata:
            lines.append(
                f"  // Generated by TreeSitter Chunker - {metadata.created_at}",
            )
            lines.append(
                f"  // Chunks: {metadata.chunk_count}, Relationships: {metadata.relationship_count}",
            )
        lines.append("")
        lines.append("  // Nodes")
        for chunk in chunks:
            label_parts = [f"{chunk.node_type}"]
            for attr in self._node_attributes:
                if attr not in {"node_type"} and hasattr(chunk, attr):
                    value = getattr(chunk, attr)
                    if value:
                        label_parts.append(f"{attr}: {value}")
            label = "\\n".join(label_parts)
            lines.append(f'  "{chunk.chunk_id}" [label="{label}"];')
        lines.append("")
        lines.append("  // Edges")
        for rel in relationships:
            edge_style = self._edge_style.get(
                rel.relationship_type.value,
                {"color": "black", "style": "solid"},
            )
            style_attrs = ", ".join(f'{k}="{v}"' for k, v in edge_style.items())
            label = rel.relationship_type.value
            if rel.metadata:
                for key in self._edge_attributes[1:]:
                    if key in rel.metadata:
                        label += f"\\n{key}: {rel.metadata[key]}"
            lines.append(
                f'  "{rel.source_chunk_id}" -> "{rel.target_chunk_id}" [label="{label}", {style_attrs}];',
            )
        lines.append("}")
        return "\n".join(lines)

    def _stream_dot(
        self,
        chunk_iterator: Iterator[CodeChunk],
        relationship_iterator: Iterator[ChunkRelationship],
        output: io.IOBase,
    ) -> None:
        """Stream DOT content to output."""
        output.write("digraph CodeStructure {\n")
        output.write(f"  rankdir={self._layout_algorithm.upper()};\n")
        output.write(
            "  node ["
            + ", ".join(f'{k}="{v}"' for k, v in self._node_style.items())
            + "];\n",
        )
        output.write("\n")
        output.write("  // Nodes\n")
        for chunk in chunk_iterator:
            label_parts = [f"{chunk.node_type}"]
            for attr in self._node_attributes:
                if attr not in {"node_type"} and hasattr(chunk, attr):
                    value = getattr(chunk, attr)
                    if value:
                        label_parts.append(f"{attr}: {value}")
            label = "\\n".join(label_parts)
            output.write(f'  "{chunk.chunk_id}" [label="{label}"];\n')
        output.write("\n")
        output.write("  // Edges\n")
        for rel in relationship_iterator:
            edge_style = self._edge_style.get(
                rel.relationship_type.value,
                {"color": "black", "style": "solid"},
            )
            style_attrs = ", ".join(f'{k}="{v}"' for k, v in edge_style.items())
            label = rel.relationship_type.value
            if rel.metadata:
                for key in self._edge_attributes[1:]:
                    if key in rel.metadata:
                        label += f"\\n{key}: {rel.metadata[key]}"
            output.write(
                f"""  "{rel.source_chunk_id}" -> "{rel.target_chunk_id}" [label="{label}", {style_attrs}];
""",
            )
        output.write("}\n")
