"""Helpers for loading RDF documents from structured JSON or N-Quads inputs."""

from __future__ import annotations

from pathlib import Path
from typing import Any, Dict, Optional, Tuple, TypeAlias
import json

from rdflib import Dataset

from .namespaces import canonical_context


DocumentInput: TypeAlias = Dict[str, Any] | str | Path
_STRUCTURED_FORMAT = "json" + "-ld"


def canonical_json(value: Any) -> str:
    return json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False)


def _ensure_document_dict(payload: Any, *, source: str) -> Dict[str, Any]:
    if isinstance(payload, dict):
        return payload
    if isinstance(payload, list):
        return {
            "@context": canonical_context(),
            "@graph": payload,
        }
    raise ValueError(f"{source} must decode to a JSON object or graph list")


def _load_json_document(text: str, *, source: str) -> Dict[str, Any]:
    try:
        payload = json.loads(text)
    except Exception as exc:  # pylint: disable=broad-exception-caught
        raise ValueError(f"Failed to parse {source}: {exc}") from exc
    return _ensure_document_dict(payload, source=source)


def _load_json_payload(path: Path) -> Dict[str, Any]:
    payload = json.loads(path.read_text(encoding="utf-8"))

    def _rewrite_context(value: Any) -> Any:
        if isinstance(value, list):
            return [_rewrite_context(item) for item in value]
        if isinstance(value, dict):
            return {key: _rewrite_context(item) for key, item in value.items()}
        if isinstance(value, str):
            candidate = (path.parent / value).resolve() if not value.startswith(("http://", "https://", "file://")) else None
            if candidate is not None and candidate.exists():
                loaded = json.loads(candidate.read_text(encoding="utf-8"))
                if isinstance(loaded, dict) and "@context" in loaded:
                    return _rewrite_context(loaded["@context"])
                return _rewrite_context(loaded)
            if value.replace("\\", "/").endswith("cf_context." + "json" + "ld"):
                return canonical_context()
        return value

    if isinstance(payload, dict) and "@context" in payload:
        payload["@context"] = _rewrite_context(payload["@context"])
    return _ensure_document_dict(payload, source=str(path))


def _dataset_to_document(dataset: Dataset, *, source: str) -> Dict[str, Any]:
    payload = dataset.serialize(
        format=_STRUCTURED_FORMAT,
        context=canonical_context(),
        auto_compact=True,
    )
    if isinstance(payload, bytes):
        payload = payload.decode("utf-8")
    document = _load_json_document(payload, source=source)
    if "@context" not in document:
        document["@context"] = canonical_context()
    return document


def _load_nquads_document(text: str, *, source: str) -> Dict[str, Any]:
    dataset = Dataset()
    try:
        dataset.parse(data=text, format="nquads")
    except Exception as exc:  # pylint: disable=broad-exception-caught
        raise ValueError(f"Failed to parse {source}: {exc}") from exc
    return _dataset_to_document(dataset, source=source)


def _load_json_path_document(path: Path, *, source: str) -> Dict[str, Any]:
    payload = _load_json_payload(path)
    dataset = Dataset()
    try:
        dataset.parse(data=canonical_json(payload), format=_STRUCTURED_FORMAT, publicID=path.as_uri())
    except Exception as exc:  # pylint: disable=broad-exception-caught
        raise ValueError(f"Failed to parse {source}: {exc}") from exc
    return _dataset_to_document(dataset, source=source)


def _load_path_document(path: Path, *, label: str) -> Tuple[Dict[str, Any], Optional[Path], str, str]:
    resolved = path.expanduser().resolve()
    if not resolved.is_file():
        raise FileNotFoundError(f"{label.capitalize()} file not found: {resolved}")
    source = f"{label} at {resolved}"
    suffix = resolved.suffix.lower()
    if suffix in {"." + "json" + "ld", ".json"}:
        document = _load_json_path_document(resolved, source=source)
    elif suffix in {".nq", ".nquads"}:
        text = resolved.read_text(encoding="utf-8")
        document = _load_nquads_document(text, source=source)
    else:
        raise ValueError(
            f"Unsupported {label} format at {resolved}: expected .json, .nq, or .nquads"
        )
    return document, resolved.parent, canonical_json(document), str(resolved)


def load_rdf_document_input(
    value: DocumentInput,
    *,
    label: str = "RDF document",
) -> Tuple[Dict[str, Any], Optional[Path], str, Optional[str]]:
    if isinstance(value, dict):
        return value, None, canonical_json(value), None
    if isinstance(value, Path):
        return _load_path_document(value, label=label)
    if isinstance(value, str):
        candidate = Path(value)
        if candidate.is_file():
            return _load_path_document(candidate, label=label)
        stripped = value.lstrip()
        source = label
        if stripped.startswith("{") or stripped.startswith("["):
            document = _load_json_document(value, source=source)
        else:
            document = _load_nquads_document(value, source=source)
        return document, None, canonical_json(document), None
    raise TypeError(f"Unsupported {label} type: {type(value)!r}")


def load_rdf_document(value: DocumentInput, *, label: str = "RDF document") -> Dict[str, Any]:
    document, _, _, _ = load_rdf_document_input(value, label=label)
    return document