"""Canonical RDFLib Dataset/Graph view helpers for DuckDB-backed semantics."""

from __future__ import annotations

from typing import Any, Dict, Iterable, Optional
import json

from rdflib import Dataset, Graph

from .namespaces import canonical_context


_STRUCTURED_FORMAT = "json" + "-ld"
_DEFAULT_GRAPH_ID = "urn:x-rdflib:default"


def canonical_json(value: Any) -> str:
    """Serialize JSON deterministically for RDF parsing/exports."""
    return json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False)


def merge_contexts(*contexts: Optional[Dict[str, Any]]) -> Dict[str, Any]:
    """Merge RDF document contexts with canonical CogniFlow prefixes."""
    merged: Dict[str, Any] = dict(canonical_context())
    for context in contexts:
        if not isinstance(context, dict):
            continue
        for key, value in context.items():
            merged[key] = value
    return merged


def dataset_from_document(
    document: Dict[str, Any],
    *,
    public_id: Optional[str] = None,
) -> Dataset:
    """Parse a JSON-LD style RDF document into an RDFLib Dataset."""
    dataset = Dataset()
    parse_kwargs = {
        "data": canonical_json(document),
        "format": _STRUCTURED_FORMAT,
    }
    if public_id is not None:
        parse_kwargs["publicID"] = public_id
    dataset.parse(**parse_kwargs)
    return dataset


def graph_from_document(
    document: Dict[str, Any],
    *,
    public_id: Optional[str] = None,
) -> Graph:
    """Parse a JSON-LD style RDF document into a flattened RDFLib Graph."""
    return union_graph(dataset_from_document(document, public_id=public_id))


def append_graph(target: Graph, source: Graph) -> Graph:
    """Copy triples from one graph into another graph."""
    for subj, pred, obj in source:
        target.add((subj, pred, obj))
    return target


def append_dataset(target: Dataset, source: Dataset) -> Dataset:
    """Copy quads from one dataset into another dataset."""
    for graph in iter_graphs(source):
        sink = target.graph(graph.identifier)
        for subj, pred, obj in graph:
            sink.add((subj, pred, obj))
    return target


def iter_graphs(dataset: Dataset):
    """Iterate dataset graphs while skipping empty RDFLib default graph shells."""
    graphs = dataset.graphs() if hasattr(dataset, "graphs") else dataset.contexts()
    for graph in graphs:
        if str(graph.identifier) == _DEFAULT_GRAPH_ID and len(graph) == 0:
            continue
        yield graph


def union_graph(dataset: Dataset) -> Graph:
    """Return a flattened Graph view across all named graphs in a Dataset."""
    graph = Graph()
    for named_graph in iter_graphs(dataset):
        append_graph(graph, named_graph)
    return graph


def dataset_to_document(
    dataset: Dataset,
    *,
    context: Optional[Dict[str, Any]] = None,
) -> Dict[str, Any]:
    """Serialize a Dataset into a JSON-LD style dataset document."""
    merged_context = merge_contexts(context)
    payload = dataset.serialize(
        format=_STRUCTURED_FORMAT,
        context=merged_context,
        auto_compact=True,
    )
    if isinstance(payload, bytes):
        payload = payload.decode("utf-8")
    document = json.loads(payload)
    if isinstance(document, list):
        return {"@context": merged_context, "@graph": document}
    if isinstance(document, dict):
        if "@graph" in document:
            document.setdefault("@context", merged_context)
            return document
        return {"@context": merged_context, "@graph": [document]}
    raise ValueError("Unsupported JSON-LD payload produced from RDFLib Dataset")


def graph_to_document(
    graph: Graph,
    *,
    context: Optional[Dict[str, Any]] = None,
) -> Dict[str, Any]:
    """Serialize a Graph into a JSON-LD style graph document."""
    merged_context = merge_contexts(context)
    payload = graph.serialize(
        format=_STRUCTURED_FORMAT,
        context=merged_context,
        auto_compact=True,
    )
    if isinstance(payload, bytes):
        payload = payload.decode("utf-8")
    document = json.loads(payload)
    if isinstance(document, list):
        return {"@context": merged_context, "@graph": document}
    if isinstance(document, dict):
        if "@graph" in document:
            document.setdefault("@context", merged_context)
            return document
        return {"@context": merged_context, "@graph": [document]}
    raise ValueError("Unsupported JSON-LD payload produced from RDFLib Graph")
