"""Export knowledge graphs to various formats.

Supports GraphML, GEXF (Gephi), and CSV. All formats flatten
complex attributes (lists, dicts) to strings for compatibility.
"""

import csv
import json
import logging
from pathlib import Path
from typing import Any

import networkx as nx

from sift_kg.graph.knowledge_graph import KnowledgeGraph
from sift_kg.graph.postprocessor import strip_metadata
from sift_kg.visualize import DEFAULT_ENTITY_COLOR, EDGE_PALETTE, ENTITY_COLORS

logger = logging.getLogger(__name__)


def _hex_to_rgb(hex_color: str) -> tuple[int, int, int]:
    """Convert hex color to (r, g, b) tuple."""
    h = hex_color.lstrip("#")
    return int(h[0:2], 16), int(h[2:4], 16), int(h[4:6], 16)

SUPPORTED_FORMATS = ("json", "graphml", "gexf", "csv")


def export_graph(
    kg: KnowledgeGraph,
    output_path: Path,
    fmt: str,
    descriptions: dict[str, str] | None = None,
) -> Path:
    """Export a knowledge graph to the specified format.

    Args:
        kg: Knowledge graph to export
        output_path: Where to write the output (file or directory for CSV)
        fmt: Format string — "json", "graphml", "gexf", or "csv"
        descriptions: Optional entity descriptions to merge into node attributes

    Returns:
        Path to the written file (or directory for CSV)
    """
    fmt = fmt.lower()
    if fmt not in SUPPORTED_FORMATS:
        raise ValueError(f"Unsupported format: {fmt}. Supported: {', '.join(SUPPORTED_FORMATS)}")

    output_path.parent.mkdir(parents=True, exist_ok=True)

    if fmt == "json":
        return _export_json(kg, output_path)

    # Strip DOCUMENT nodes + MENTIONED_IN edges for visual/analytical exports
    clean = strip_metadata(kg)

    if fmt == "graphml":
        return _export_graphml(clean, output_path, descriptions)
    elif fmt == "gexf":
        return _export_gexf(clean, output_path, descriptions)
    elif fmt == "csv":
        return _export_csv(clean, output_path, descriptions)
    raise ValueError(f"Unsupported format: {fmt}")


def _export_json(kg: KnowledgeGraph, output_path: Path) -> Path:
    """Export as sift-kg native JSON."""
    kg.save(output_path)
    return output_path


def _flatten_value(value: Any) -> str | int | float | bool:
    """Flatten complex values to strings for GraphML/GEXF compatibility."""
    if isinstance(value, (str, int, float, bool)):
        return value
    if isinstance(value, list):
        return "; ".join(str(v) for v in value)
    if isinstance(value, dict):
        return json.dumps(value, default=str)
    return str(value)


def _build_flat_graph(
    kg: KnowledgeGraph,
    descriptions: dict[str, str] | None = None,
) -> nx.DiGraph:
    """Build a simple DiGraph with flattened attributes.

    GraphML/GEXF don't support MultiDiGraph well (parallel edges get
    collapsed). We merge parallel edges by concatenating relation types.

    Sets `label` on nodes and edges for Gephi/yEd display.
    Optionally merges entity descriptions from narrate output.
    """
    flat = nx.DiGraph()
    descriptions = descriptions or {}
    degrees = dict(kg.graph.degree())

    for node_id, data in kg.graph.nodes(data=True):
        flat_attrs = {k: _flatten_value(v) for k, v in data.items()}
        # Gephi uses 'label' for display — set it to entity name
        flat_attrs["label"] = data.get("name", node_id)
        # Color by entity type — same palette as pyvis viewer
        entity_type = data.get("entity_type", "UNKNOWN")
        color = ENTITY_COLORS.get(entity_type, DEFAULT_ENTITY_COLOR)
        r, g, b = _hex_to_rgb(color)
        flat_attrs["color"] = color
        flat_attrs["r"] = r
        flat_attrs["g"] = g
        flat_attrs["b"] = b
        # Size by degree — same formula as pyvis viewer
        degree = degrees.get(node_id, 0)
        flat_attrs["size"] = float(max(8, min(50, 6 + degree * 2.5)))
        # Merge description if available
        if node_id in descriptions:
            flat_attrs["description"] = descriptions[node_id]
        flat.add_node(node_id, **flat_attrs)

    # Merge parallel edges between same source/target
    rel_color_map: dict[str, str] = {}
    edge_map: dict[tuple[str, str], dict[str, Any]] = {}
    for source, target, _key, data in kg.graph.edges(data=True, keys=True):
        pair = (source, target)
        if pair not in edge_map:
            edge_map[pair] = {k: _flatten_value(v) for k, v in data.items()}
            # Track multi-valued fields as sets for proper merging
            edge_map[pair]["_relation_types"] = {data.get("relation_type", "")}
            edge_map[pair]["_evidences"] = {data.get("evidence", "")} - {""}
        else:
            # Merge relation types (set-based, no substring issues)
            edge_map[pair]["_relation_types"].add(data.get("relation_type", ""))
            edge_map[pair]["_evidences"].add(data.get("evidence", ""))
            # Keep highest confidence
            new_conf = data.get("confidence", 0)
            if isinstance(new_conf, (int, float)) and new_conf > edge_map[pair].get("confidence", 0):
                edge_map[pair]["confidence"] = new_conf

    # Flatten merged sets back to strings + assign colors
    for attrs in edge_map.values():
        types = attrs.pop("_relation_types", set())
        rel_type = "; ".join(sorted(t for t in types if t))
        attrs["relation_type"] = rel_type
        evidences = attrs.pop("_evidences", set())
        if evidences:
            attrs["evidence"] = " | ".join(sorted(evidences))
        # Color by relation type — same palette as pyvis viewer
        primary_type = next((t for t in sorted(types) if t), rel_type)
        if primary_type not in rel_color_map:
            idx = len(rel_color_map) % len(EDGE_PALETTE)
            rel_color_map[primary_type] = EDGE_PALETTE[idx]
        color = rel_color_map[primary_type]
        r, g, b = _hex_to_rgb(color)
        attrs["color"] = color
        attrs["r"] = r
        attrs["g"] = g
        attrs["b"] = b

    for (source, target), attrs in edge_map.items():
        # Gephi uses 'label' for edge display
        attrs["label"] = attrs.get("relation_type", "")
        flat.add_edge(source, target, **attrs)

    # Pre-compute layout so Gephi/yEd render a spread-out graph on import
    pos = nx.spring_layout(flat, k=3.0, iterations=100, seed=42, scale=1000)
    for node_id, (x, y) in pos.items():
        flat.nodes[node_id]["x"] = float(x)
        flat.nodes[node_id]["y"] = float(y)

    return flat


def _export_graphml(
    kg: KnowledgeGraph, output_path: Path, descriptions: dict[str, str] | None = None,
) -> Path:
    """Export as GraphML (compatible with yEd, Gephi, Cytoscape)."""
    flat = _build_flat_graph(kg, descriptions)
    nx.write_graphml(flat, str(output_path))
    logger.info(f"GraphML exported: {flat.number_of_nodes()} nodes, {flat.number_of_edges()} edges -> {output_path}")
    return output_path


def _export_gexf(
    kg: KnowledgeGraph, output_path: Path, descriptions: dict[str, str] | None = None,
) -> Path:
    """Export as GEXF (Gephi native format)."""
    flat = _build_flat_graph(kg, descriptions)
    nx.write_gexf(flat, str(output_path))
    logger.info(f"GEXF exported: {flat.number_of_nodes()} nodes, {flat.number_of_edges()} edges -> {output_path}")
    return output_path


def _export_csv(
    kg: KnowledgeGraph, output_dir: Path, descriptions: dict[str, str] | None = None,
) -> Path:
    """Export as CSV (entities.csv + relations.csv)."""
    output_dir.mkdir(parents=True, exist_ok=True)
    descriptions = descriptions or {}

    # Entities
    entities_path = output_dir / "entities.csv"
    entity_rows = []
    for node_id, data in kg.graph.nodes(data=True):
        entity_rows.append({
            "id": node_id,
            "name": data.get("name", ""),
            "entity_type": data.get("entity_type", ""),
            "confidence": data.get("confidence", ""),
            "source_documents": "; ".join(data.get("source_documents", [])),
            "attributes": json.dumps(data.get("attributes", {}), default=str),
            "description": descriptions.get(node_id, ""),
        })

    entity_fields = ["id", "name", "entity_type", "confidence", "source_documents", "attributes", "description"]
    with open(entities_path, "w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=entity_fields)
        writer.writeheader()
        writer.writerows(entity_rows)

    # Relations
    relations_path = output_dir / "relations.csv"
    relation_rows = []
    for source, target, _key, data in kg.graph.edges(data=True, keys=True):
        relation_rows.append({
            "source": source,
            "target": target,
            "relation_type": data.get("relation_type", ""),
            "confidence": data.get("confidence", ""),
            "evidence": data.get("evidence", ""),
            "source_document": data.get("source_document", ""),
        })

    relation_fields = ["source", "target", "relation_type", "confidence", "evidence", "source_document"]
    with open(relations_path, "w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=relation_fields)
        writer.writeheader()
        writer.writerows(relation_rows)

    logger.info(
        f"CSV exported: {len(entity_rows)} entities, {len(relation_rows)} relations -> {output_dir}"
    )
    return output_dir
