"""Scene graph extraction and serialization helpers.

This module provides utilities to transform lightweight text metadata into a
symbolic scene graph represented as a :class:`networkx.DiGraph`.  The resulting
structure is compatible with :class:`hieromem_core.episode.Episode` and can be
persisted using the provided serializers.
"""

from __future__ import annotations

from collections.abc import Iterable
from typing import Any

import networkx as nx

__all__ = [
    "extract_scene_graph",
    "graph_to_metadata",
    "scene_graph_from_metadata",
]


def extract_scene_graph(metadata: Any) -> nx.DiGraph:
    """Build a scene graph from semi-structured metadata.

    Parameters
    ----------
    metadata:
        A dictionary, string, or iterable describing the entities and relations
        present in a scene.  Supported formats are intentionally lightweight so
        that LLM generated metadata can be consumed without an additional
        parsing step:

        - ``dict`` with ``"entities"`` and ``"relations"`` entries.
        - ``str`` containing sections labelled ``Entities:`` and
          ``Relations:`` with bullet items.
        - Any iterable yielding entity/edge descriptors that follow the above
          conventions.

    Returns
    -------
    networkx.DiGraph
        A directed acyclic graph describing the symbolic structure of the
        scene.  Edges that would introduce a cycle are ignored so the result is
        always compatible with :class:`hieromem_core.episode.Episode`.
    """

    if isinstance(metadata, nx.DiGraph):
        return metadata.copy(as_view=False)

    structured = _normalise_metadata(metadata)

    graph = nx.DiGraph()

    for entity in structured.get("entities", []):
        node_id, attrs = _parse_entity(entity)
        if node_id not in graph:
            graph.add_node(node_id, **attrs)
        else:
            graph.nodes[node_id].update(attrs)

    for relation in structured.get("relations", []):
        edge_data = _parse_relation(relation)
        if edge_data is None:
            continue
        u, v, attrs = edge_data
        if u == v:
            continue
        if u not in graph:
            graph.add_node(u)
        if v not in graph:
            graph.add_node(v)
        graph.add_edge(u, v, **attrs)
        if not nx.is_directed_acyclic_graph(graph):
            graph.remove_edge(u, v)

    return graph


def graph_to_metadata(graph: nx.DiGraph) -> dict[str, list[dict[str, Any]]]:
    """Serialise a scene graph into a JSON friendly structure."""

    nodes: list[dict[str, Any]] = []
    for node_id, attrs in graph.nodes(data=True):
        node_entry: dict[str, Any] = {"id": node_id}
        node_entry.update(attrs)
        nodes.append(node_entry)

    edges: list[dict[str, Any]] = []
    for u, v, attrs in graph.edges(data=True):
        edge_entry: dict[str, Any] = {"u": u, "v": v}
        edge_entry.update(attrs)
        edges.append(edge_entry)

    return {"nodes": nodes, "edges": edges}


def scene_graph_from_metadata(
        data: dict[str, Iterable[dict[str, Any]]]) -> nx.DiGraph:
    """Reconstruct a graph previously produced by :func:`graph_to_metadata`."""

    graph = nx.DiGraph()

    for node in data.get("nodes", []):
        node = dict(node)
        node_id = node.pop("id")
        graph.add_node(node_id, **node)

    for edge in data.get("edges", []):
        edge = dict(edge)
        u = edge.pop("u")
        v = edge.pop("v")
        graph.add_edge(u, v, **edge)

    return graph


def _normalise_metadata(metadata: Any) -> dict[str, list[Any]]:
    if metadata is None:
        return {"entities": [], "relations": []}

    if isinstance(metadata, dict):
        dict_entities = list(metadata.get("entities", []))
        dict_relations = list(metadata.get("relations", []))
        return {"entities": dict_entities, "relations": dict_relations}

    if isinstance(metadata, str):
        return _parse_text_metadata(metadata)

    if isinstance(metadata, Iterable):
        iter_entities: list[Any] = []
        iter_relations: list[Any] = []
        for item in metadata:
            if isinstance(item, dict) and {"source", "target"} <= item.keys():
                iter_relations.append(item)
            else:
                iter_entities.append(item)
        return {"entities": iter_entities, "relations": iter_relations}

    raise TypeError("Unsupported metadata format for scene graph extraction.")


def _parse_text_metadata(text: str) -> dict[str, list[str]]:
    entities: list[str] = []
    relations: list[str] = []
    current: str | None = None

    for raw_line in text.splitlines():
        line = raw_line.strip()
        if not line:
            continue
        lower = line.lower()
        if lower.startswith("entities"):
            current = "entities"
            continue
        if lower.startswith("relations"):
            current = "relations"
            continue

        if line.startswith("-") or line.startswith("*"):
            line = line[1:].strip()

        if current == "entities":
            entities.append(line)
        elif current == "relations":
            relations.append(line)

    return {"entities": entities, "relations": relations}


def _parse_entity(entity: Any) -> tuple[str, dict[str, Any]]:
    if isinstance(entity, dict):
        if "id" not in entity:
            raise ValueError("Entity dictionaries must include an 'id' field.")
        entity = dict(entity)
        node_id = str(entity.pop("id"))
        return node_id, {k: _coerce_value(v) for k, v in entity.items()}

    if not isinstance(entity, str):
        raise TypeError(f"Unsupported entity description: {entity!r}")

    parts = entity.split()
    if not parts:
        raise ValueError("Empty entity description encountered.")

    head = parts[0]
    attrs: dict[str, Any] = {}

    if ":" in head:
        node_id, node_type = head.split(":", 1)
        attrs["type"] = _coerce_value(node_type)
    else:
        node_id = head

    for token in parts[1:]:
        if "=" not in token:
            continue
        key, value = token.split("=", 1)
        attrs[key] = _coerce_value(value)

    return node_id, attrs


def _parse_relation(relation: Any) -> tuple[str, str, dict[str, Any]] | None:
    if isinstance(relation, dict):
        data = dict(relation)
        source = data.pop("source", data.pop("u", None))
        target = data.pop("target", data.pop("v", None))
        if source is None or target is None:
            raise ValueError(
                "Relation dictionaries require 'source'/'target' keys.")
        return str(source), str(target), {
            k: _coerce_value(v) for k, v in data.items()}

    if not isinstance(relation, str):
        raise TypeError(f"Unsupported relation description: {relation!r}")

    body, _, trailing = relation.partition("|")
    body = body.strip()
    trailing = trailing.strip()

    if "->" not in body:
        return None

    source_part, target_part = body.split("->", 1)
    source = source_part.strip()
    target = target_part.strip()

    attrs: dict[str, Any] = {}
    if trailing:
        tokens = trailing.split()
        relation_set = False
        for index, token in enumerate(tokens):
            if "=" in token:
                key, value = token.split("=", 1)
                if key == "relation":
                    attrs["relation"] = _coerce_value(value)
                    relation_set = True
                else:
                    attrs[key] = _coerce_value(value)
                continue

            if not relation_set and index == 0:
                attrs["relation"] = _coerce_value(token)
                relation_set = True

    return source, target, attrs


def _coerce_value(value: Any) -> Any:
    if isinstance(value, int | float | bool):
        return value

    if not isinstance(value, str):
        return value

    lower = value.lower()
    if lower in {"true", "false"}:
        return lower == "true"

    try:
        if "." in value:
            return float(value)
        return int(value)
    except ValueError:
        return value
