"""Small canonical enums and their normalisers. Raw source strings are never mutated in snapshots — only the materialised attribute
is canonical; connectors and the writer call these before writing a claim."""
from __future__ import annotations

import re
from typing import Any

# ---------------------------------------------------------------------------------------------- modalities
MODALITIES = ("text", "image", "audio", "video", "document", "code", "embedding", "3d", "structured", "action")
_MODALITY_ALIASES = {
    "text": "text", "texts": "text", "language": "text", "nlp": "text", "chat": "text", "textual": "text",
    "image": "image", "images": "image", "vision": "image", "visual": "image", "picture": "image", "photo": "image", "img": "image",
    "audio": "audio", "speech": "audio", "voice": "audio", "sound": "audio", "music": "audio",
    "video": "video", "videos": "video",
    "pdf": "document", "document": "document", "documents": "document", "file": "document", "files": "document", "doc": "document",
    "code": "code", "coding": "code", "programming": "code",
    "embedding": "embedding", "embeddings": "embedding", "vector": "embedding",
    "3d": "3d", "mesh": "3d", "point cloud": "3d",
    "structured": "structured", "json": "structured", "table": "structured", "tabular": "structured", "time-series": "structured", "time series": "structured",
    "action": "action", "actions": "action", "computer-use": "action", "computer use": "action", "tool": "action", "robotics": "action",
    "multimodal": None, "multi-modal": None, "omni": None, "any": None,  # not a modality — a property of the set
}


def normalize_modality(raw: Any) -> str | None:
    if raw is None:
        return None
    s = str(raw).strip().lower().replace("_", "-")
    if s in _MODALITY_ALIASES:
        return _MODALITY_ALIASES[s]
    s2 = s.replace("-", " ")
    return _MODALITY_ALIASES.get(s2)


def normalize_modalities(raw: Any) -> list[str]:
    """Any iterable / comma string → sorted canonical unique list (unknown values dropped)."""
    if raw is None:
        return []
    items = raw if isinstance(raw, (list, tuple, set)) else re.split(r"[,/;+&]|\band\b|→|->", str(raw))
    out: set[str] = set()
    for it in items:
        m = normalize_modality(it)
        if m:
            out.add(m)
    return sorted(out)


# ---------------------------------------------------------------------------------------------- model / entity status
MODEL_STATUSES = ("announced", "preview", "active", "limited-availability", "deprecated", "retired", "archived", "unknown")
_STATUS_ALIASES = {
    "active": "active", "available": "active", "ga": "active", "general availability": "active", "generally available": "active", "live": "active",
    "released": "active", "stable": "active", "current": "active", "supported": "active", "production": "active", "online": "active",
    "preview": "preview", "beta": "preview", "alpha": "preview", "experimental": "preview", "public preview": "preview", "early access": "preview",
    "research preview": "preview", "exp": "preview",
    "limited-availability": "limited-availability", "limited availability": "limited-availability", "limited": "limited-availability",
    "waitlist": "limited-availability", "private preview": "limited-availability", "invite-only": "limited-availability", "gated": "limited-availability",
    "deprecated": "deprecated", "legacy": "deprecated", "sunset": "deprecated", "sunsetting": "deprecated", "end-of-life announced": "deprecated",
    "deprecation": "deprecated",
    "retired": "retired", "discontinued": "retired", "shut down": "retired", "shutdown": "retired", "removed": "retired", "end of life": "retired",
    "eol": "retired", "decommissioned": "retired", "offline": "retired", "delisted": "retired",
    "announced": "announced", "upcoming": "announced", "coming soon": "announced", "planned": "announced", "unreleased": "announced",
    "archived": "archived",
    "unknown": "unknown", "": "unknown",
}


def normalize_status(raw: Any) -> str | None:
    if raw is None:
        return None
    s = str(raw).strip().lower().replace("_", "-")
    if s in _STATUS_ALIASES:
        return _STATUS_ALIASES[s]
    return _STATUS_ALIASES.get(s.replace("-", " "))


# ---------------------------------------------------------------------------------------------- organization kinds
ORG_KINDS = ("company", "lab", "university", "nonprofit", "government", "community", "consortium", "individual")
_ORG_KIND_ALIASES = {
    "company": "company", "corporation": "company", "startup": "company", "enterprise": "company", "vendor": "company", "business": "company",
    "lab": "lab", "laboratory": "lab", "research lab": "lab", "research-lab": "lab", "research institute": "lab", "institute": "lab", "research": "lab",
    "university": "university", "academic": "university", "school": "university", "college": "university",
    "nonprofit": "nonprofit", "non-profit": "nonprofit", "non profit": "nonprofit", "foundation": "nonprofit", "ngo": "nonprofit",
    "government": "government", "public": "government", "agency": "government", "state": "government",
    "community": "community", "open-source community": "community", "collective": "community", "hf-community": "community", "group": "community",
    "consortium": "consortium", "alliance": "consortium", "coalition": "consortium",
    "individual": "individual", "person": "individual", "independent": "individual", "user": "individual",
}
# entity_type → default org kind (the entity_type stays; org_kind is the canonical sub-kind)
ORG_TYPE_DEFAULT_KIND = {"company": "company", "lab": "lab", "university": "university", "organization": None}


def normalize_org_kind(raw: Any) -> str | None:
    if raw is None:
        return None
    s = str(raw).strip().lower().replace("_", "-")
    return _ORG_KIND_ALIASES.get(s) or _ORG_KIND_ALIASES.get(s.replace("-", " "))


# ---------------------------------------------------------------------------------------------- hardware kinds
HARDWARE_KINDS = ("gpu", "accelerator", "npu", "cpu", "soc", "system", "server", "workstation", "cloud-instance", "rack")
_HW_ALIASES = {
    "gpu": "gpu", "graphics card": "gpu", "graphics": "gpu", "dgpu": "gpu",
    "accelerator": "accelerator", "tpu": "accelerator", "asic": "accelerator", "ai accelerator": "accelerator", "lpu": "accelerator", "wse": "accelerator",
    "npu": "npu", "neural engine": "npu",
    "cpu": "cpu", "processor": "cpu",
    "soc": "soc", "system on chip": "soc", "system-on-chip": "soc", "apple silicon": "soc", "chip": "soc",
    "system": "system", "computer": "system", "desktop": "system", "laptop": "system", "mini": "system", "mac": "system",
    "server": "server", "node": "server", "dgx": "server", "appliance": "server",
    "workstation": "workstation",
    "cloud-instance": "cloud-instance", "cloud instance": "cloud-instance", "instance": "cloud-instance", "vm": "cloud-instance",
    "rack": "rack", "superpod": "rack", "pod": "rack", "cluster": "rack", "nvl72": "rack",
}


def normalize_hardware_kind(raw: Any) -> str | None:
    if raw is None:
        return None
    s = str(raw).strip().lower().replace("_", "-")
    return _HW_ALIASES.get(s) or _HW_ALIASES.get(s.replace("-", " "))


# ---------------------------------------------------------------------------------------------- framework / library / tool kinds
FRAMEWORK_KINDS = ("training-framework", "inference-engine", "serving-engine", "library", "runtime", "agent-framework", "orchestration",
                   "evaluation-harness", "sdk", "tool", "application", "agent", "mcp-server", "vector-database", "observability", "data-tooling")
_FW_ALIASES = {
    "framework": "training-framework", "training framework": "training-framework", "training-framework": "training-framework", "deep learning framework": "training-framework",
    "inference engine": "inference-engine", "inference-engine": "inference-engine", "inference": "inference-engine", "engine": "inference-engine",
    "serving": "serving-engine", "serving engine": "serving-engine", "serving-engine": "serving-engine", "server": "serving-engine",
    "library": "library", "lib": "library", "package": "library",
    "runtime": "runtime",
    "agent framework": "agent-framework", "agent-framework": "agent-framework", "agents": "agent-framework",
    "orchestration": "orchestration", "workflow": "orchestration", "pipeline": "orchestration",
    "evaluation": "evaluation-harness", "eval": "evaluation-harness", "evaluation harness": "evaluation-harness", "evaluation-harness": "evaluation-harness", "benchmarking": "evaluation-harness",
    "sdk": "sdk", "client": "sdk", "api client": "sdk",
    "tool": "tool", "cli": "tool", "utility": "tool",
    "application": "application", "app": "application", "product": "application", "ui": "application", "desktop app": "application",
    "agent": "agent", "coding agent": "agent", "assistant": "agent",
    "mcp": "mcp-server", "mcp server": "mcp-server", "mcp-server": "mcp-server", "mcp_server": "mcp-server",
    "vector database": "vector-database", "vector-database": "vector-database", "vector db": "vector-database", "vector store": "vector-database",
    "observability": "observability", "tracing": "observability", "monitoring": "observability", "experiment tracking": "observability",
    "data": "data-tooling", "data tooling": "data-tooling", "data-tooling": "data-tooling", "datasets": "data-tooling", "tokenizer": "data-tooling",
}


def normalize_framework_kind(raw: Any) -> str | None:
    if raw is None:
        return None
    s = str(raw).strip().lower()
    return _FW_ALIASES.get(s) or _FW_ALIASES.get(s.replace("_", "-")) or _FW_ALIASES.get(s.replace("-", " "))


# ---------------------------------------------------------------------------------------------- property-level normalisation (writer + canonicalizer)
# property → taxonomy domain (the entity type picks the normaliser for `kind`)
TAXONOMY_PROPERTIES = {"license", "openness", "modalities", "modalities_input", "modalities_output", "status", "kind", "org_kind"}
_ORG_TYPES = {"company", "organization", "lab", "university", "provider"}
_FRAMEWORK_TYPES = {"framework", "library", "tool", "runtime", "agent", "mcp_server", "repository"}


def normalize_property(entity_type: str | None, prop: str, value: Any) -> tuple[Any, str | None, list[tuple[str, str, str | None]]]:
    """Canonicalise one attribute value at write time.

    Returns `(canonical_value, value_raw, mappings)` where `value_raw` is the source string when it differed from the canonical value
    (None otherwise) and `mappings` lists `(domain, raw, canonical|None)` pairs observed — unknown values are KEPT AS-IS (never dropped)
    and mapped to canonical=None so the taxonomy backlog stays visible."""
    if value is None or prop not in TAXONOMY_PROPERTIES:
        return value, None, []
    if prop == "license":
        from aiatlas.ontology.licenses import normalize_license

        return _scalar("license", value, normalize_license)
    if prop == "openness":
        from aiatlas.ontology.openness import normalize_openness

        return _scalar("openness", value, normalize_openness)
    if prop == "status":
        return _scalar("status", value, normalize_status)
    if prop == "org_kind":
        return _scalar("org_kind", value, normalize_org_kind)
    if prop == "kind":
        if entity_type == "hardware":
            return _scalar("hardware_kind", value, normalize_hardware_kind)
        if entity_type in _FRAMEWORK_TYPES:
            return _scalar("framework_kind", value, normalize_framework_kind)
        return value, None, []
    if prop in ("modalities", "modalities_input", "modalities_output"):
        items = value if isinstance(value, (list, tuple, set)) else re.split(r"[,/;+&]|\band\b|→|->", str(value))
        out: list[str] = []
        mappings: list[tuple[str, str, str | None]] = []
        changed = False
        for it in items:
            raw = str(it).strip()
            if not raw:
                continue
            canon = normalize_modality(raw)
            if canon is None and raw.lower() in _MODALITY_ALIASES:      # "multimodal" & co: a property of the set, not a modality
                mappings.append(("modality", raw, None))
                changed = True
                continue
            mappings.append(("modality", raw, canon))
            keep = canon or raw
            if keep != raw:
                changed = True
            if keep not in out:
                out.append(keep)
        canonical = sorted(out)
        if not isinstance(value, (list, tuple, set)):
            changed = True
        elif sorted(str(x).strip() for x in value if str(x).strip()) != canonical:
            changed = True
        raw_repr = (", ".join(str(x) for x in value) if isinstance(value, (list, tuple, set)) else str(value)) if changed else None
        return canonical, raw_repr, mappings
    return value, None, []


# raw values that legitimately have no canonical counterpart (kept as-is, never reported as an unmapped violation)
NOT_A_VIOLATION = {("framework_kind", "model"), ("org_kind", "organization")}


def _scalar(domain: str, value: Any, normaliser) -> tuple[Any, str | None, list[tuple[str, str, str | None]]]:
    if not isinstance(value, str):
        return value, None, []
    raw = value.strip()
    canon = normaliser(raw)
    if canon is None:
        if (domain, raw.lower()) in NOT_A_VIOLATION:
            return raw, None, []
        return raw, None, [(domain, raw, None)]
    return canon, (raw if raw != canon else None), [(domain, raw, canon)]


# ---------------------------------------------------------------------------------------------- generic helper
def canonical_enum(value: Any, normaliser) -> tuple[Any, str | None]:
    """Return (canonical_or_original, raw_if_changed). Lets the writer store `x` canonical and `x_raw` when the source label differed."""
    if value is None:
        return None, None
    canon = normaliser(value)
    if canon is None:
        return value, None
    return canon, (str(value) if str(value) != canon else None)


__all__ = [
    "FRAMEWORK_KINDS", "HARDWARE_KINDS", "MODALITIES", "MODEL_STATUSES", "ORG_KINDS", "ORG_TYPE_DEFAULT_KIND", "TAXONOMY_PROPERTIES", "canonical_enum",
    "normalize_framework_kind", "normalize_hardware_kind", "normalize_modalities", "normalize_modality", "normalize_org_kind", "normalize_property",
    "normalize_status",
]