from __future__ import annotations

import json
from pathlib import Path
from typing import Literal

import yaml
from cognite.client import CogniteClient
from rdflib import Graph
from rdflib.namespace import Namespace

from cognite_data_quality._config import DataModelConfig, RuleConfig
from cognite_data_quality._records import RecordsConfig


def load_ttl_rules(
    source: Path | str,
    *,
    client: CogniteClient | None = None,
    file_external_id: str | None = None,
) -> str:
    """Load SHACL rules as Turtle from a local file or CDF File external ID."""
    if file_external_id:
        if client is None:
            raise ValueError("client is required when loading SHACL rules by external ID")
        file_bytes = client.files.download_bytes(external_id=file_external_id)
        return file_bytes.decode("utf-8")

    path = Path(source)
    if not path.exists():
        raise FileNotFoundError(f"SHACL rules file not found: {path}")
    return path.read_text(encoding="utf-8")


def load_json_rules(source: Path | str | dict) -> str:
    """Load SHACL rules from JSON-LD rule set into Turtle."""
    if isinstance(source, str | Path):
        path = Path(source)
        if not path.exists():
            raise FileNotFoundError(f"JSON rules file not found: {path}")
        content = json.loads(path.read_text(encoding="utf-8"))
    else:
        content = source

    quality = content.get("quality", [])
    if not quality:
        raise ValueError("JSON rules content missing 'quality' list")

    graph = Graph()
    for rule_set in quality:
        rule_definitions = rule_set.get("ruleDefinitions", [])
        if not rule_definitions:
            continue
        if isinstance(rule_definitions, dict):
            rule_definitions = [rule_definitions]
        for rule_def in rule_definitions:
            graph.parse(data=json.dumps(rule_def), format="json-ld")

    return graph.serialize(format="turtle")


def load_yaml_rules(
    source: Path | str,
    *,
    client: CogniteClient | None = None,
    shacl_base_dir: Path | None = None,
) -> tuple[str, RuleConfig]:
    """Load SHACL rules from a YAML view config, resolving the TTL file path."""
    path = Path(source)
    if not path.exists():
        raise FileNotFoundError(f"YAML view config not found: {path}")

    config = yaml.safe_load(path.read_text(encoding="utf-8")) or {}
    shacl_rules = config.get("shacl_rules", {}) or {}
    shacl_file = shacl_rules.get("file")
    shacl_external_id = shacl_rules.get("external_id")

    if shacl_file is None and shacl_external_id is None:
        raise ValueError("YAML config missing shacl_rules.file or shacl_rules.external_id")

    base_dir = shacl_base_dir or path.parent
    if shacl_file:
        candidate = Path(shacl_file)
        if not candidate.is_absolute():
            candidate = base_dir / candidate
        shacl_turtle = load_ttl_rules(candidate, client=client, file_external_id=shacl_external_id)
    else:
        shacl_turtle = load_ttl_rules(path, client=client, file_external_id=shacl_external_id)

    datamodel = config.get("datamodel") or {}
    records_config = config.get("records")
    rule_config = RuleConfig(
        datamodel=DataModelConfig(
            space=datamodel.get("space", ""),
            external_id=datamodel.get("external_id", ""),
            version=datamodel.get("version", ""),
        )
        if datamodel
        else None,
        instance_space=config.get("instance_space"),
        auto_load_depth=(config.get("validation") or {}).get("auto_load_depth", 2),
        verbose=(config.get("validation") or {}).get("verbose", True),
        records=RecordsConfig(**records_config) if records_config else None,
    )

    return shacl_turtle, rule_config


def load_rules(
    source: Path | str,
    *,
    format: Literal["ttl", "json", "yaml"] | None = None,
    client: CogniteClient | None = None,
    shacl_base_dir: Path | None = None,
) -> tuple[str, RuleConfig]:
    """Load rules from TTL, JSON, or YAML into Turtle and a RuleConfig."""
    if format is None:
        suffix = Path(source).suffix.lower()
        if suffix == ".ttl":
            format = "ttl"
        elif suffix in {".json", ".jsonld"}:
            format = "json"
        elif suffix in {".yaml", ".yml"}:
            format = "yaml"
        else:
            raise ValueError(f"Could not infer rule format from suffix: {suffix}")

    if format == "ttl":
        return load_ttl_rules(source, client=client), RuleConfig()
    if format == "json":
        return load_json_rules(source), RuleConfig()
    if format == "yaml":
        return load_yaml_rules(source, client=client, shacl_base_dir=shacl_base_dir)

    raise ValueError(f"Unsupported rules format: {format}")


def extract_target_views_from_shacl(shacl_path: Path | str) -> set[str]:
    """Extract view external_ids from sh:targetClass in a SHACL TTL file.

    Parses the rules file and returns the set of view names (local names from
    sh:targetClass URIs). Use to determine which views to validate.

    Args:
        shacl_path: Path to the SHACL rules file (.ttl).

    Returns:
        Set of view external_ids (e.g. {"SmallBoat", "NMEATimeSeries"}).
    """
    sh = Namespace("http://www.w3.org/ns/shacl#")
    path = Path(shacl_path)
    if not path.exists():
        return set()
    g = Graph()
    g.parse(path, format="turtle")
    views: set[str] = set()
    for _, _, o in g.triples((None, sh.targetClass, None)):
        uri = str(o)
        local = uri.split("#")[-1] if "#" in uri else uri.split("/")[-1]
        if local:
            views.add(local)
    return views
