"""Ontology — formal schema for the knowledge graph.

Architecture:
- Engine (this file): validation, inference, loading. Stable code.
- Vocabulary (ontology.yaml): entity types, predicates, inference rules. Agent-extensible.
- Skills (SKILL.md): strategy for when/how to propose changes.

Code provides the engine. YAML provides the vocabulary. Skills provide the strategy.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field
from enum import Enum
from pathlib import Path

log = logging.getLogger(__name__)

# Authority weights for scoring — explicit facts outrank inferred ones
AUTHORITY_WEIGHTS: dict[str, float] = {"explicit": 3.0, "system": 2.0, "inferred": 1.0}


# ── Entity Type Hierarchy ──
# These stay as a Python enum because they're used in type signatures.
# The YAML can add aliases but the core types are stable.


class EntityType(str, Enum):
    """What kind of thing an entity is."""

    THING = "thing"

    # Agents — things that act
    PERSON = "person"
    ORGANIZATION = "organization"
    AI_AGENT = "ai_agent"

    # Abstracts — ideas, happenings, choices
    CONCEPT = "concept"
    SKILL = "skill"
    EVENT = "event"
    DECISION = "decision"

    # Artifacts — created things
    PROJECT = "project"
    TASK = "task"
    GOAL = "goal"
    DOCUMENT = "document"
    SOURCE_SUMMARY = "source_summary"

    # Places
    PLACE = "place"


# Groups — used by YAML domain/range references
ENTITY_GROUPS: dict[str, set[EntityType]] = {
    "agent": {EntityType.PERSON, EntityType.ORGANIZATION, EntityType.AI_AGENT},
    "abstract": {EntityType.CONCEPT, EntityType.SKILL, EntityType.EVENT, EntityType.DECISION},
    "artifact": {EntityType.PROJECT, EntityType.TASK, EntityType.GOAL, EntityType.DOCUMENT, EntityType.SOURCE_SUMMARY},
    "place": {EntityType.PLACE},
    "any": set(EntityType),
}

# Convenience aliases
AGENT_TYPES = ENTITY_GROUPS["agent"]
ABSTRACT_TYPES = ENTITY_GROUPS["abstract"]
ARTIFACT_TYPES = ENTITY_GROUPS["artifact"]
ALL_TYPES = ENTITY_GROUPS["any"]

# Legacy type mapping (string -> EntityType)
LEGACY_TYPE_MAP: dict[str, EntityType] = {
    "person": EntityType.PERSON,
    "company": EntityType.ORGANIZATION,
    "organization": EntityType.ORGANIZATION,
    "org": EntityType.ORGANIZATION,
    "institution": EntityType.ORGANIZATION,
    "project": EntityType.PROJECT,
    "concept": EntityType.CONCEPT,
    "topic": EntityType.CONCEPT,
    "idea": EntityType.CONCEPT,
    "theme": EntityType.CONCEPT,
    "place": EntityType.PLACE,
    "location": EntityType.PLACE,
    "city": EntityType.PLACE,
    "country": EntityType.PLACE,
    "region": EntityType.PLACE,
    "event": EntityType.EVENT,
    "meeting": EntityType.EVENT,
    "milestone": EntityType.EVENT,
    "task": EntityType.TASK,
    "goal": EntityType.GOAL,
    "agent": EntityType.AI_AGENT,
    "bot": EntityType.AI_AGENT,
    "skill": EntityType.SKILL,
    "entity": EntityType.THING,
    "thing": EntityType.THING,
    "document": EntityType.DOCUMENT,
    "file": EntityType.DOCUMENT,
    "article": EntityType.DOCUMENT,
    "note": EntityType.DOCUMENT,
    "decision": EntityType.DECISION,
}


def resolve_entity_type(raw: str) -> EntityType:
    """Convert a raw type string to EntityType, with fallback."""
    try:
        return EntityType(raw)
    except ValueError:
        return LEGACY_TYPE_MAP.get(raw.lower(), EntityType.THING)


def is_subtype(child: EntityType, parent: EntityType) -> bool:
    """Check if child is a subtype of parent (or equal).

    Only THING is a universal parent. Group membership (agent, artifact, etc.)
    does NOT make members subtypes of each other — person is not a subtype of
    organization, even though both are in the 'agent' group.
    """
    if child == parent or parent == EntityType.THING:
        return True
    return False


# ── Epistemic Status ──


class EpistemicStatus(str, Enum):
    """How justified is this piece of knowledge?

    Orthogonal to authority (who said it). A user can state something uncertain,
    and an agent can observe something verified.

    Phase 6 (Memory Dreaming) addition: ``REJECTED_HYPOTHESIS`` is the
    counterfactual layer — facts the agent considered and explicitly
    rejected, with the evidence that defeated them. This is *not*
    just "wrong, ignore it" — it's "the agent's intellectual history
    of considered-and-discarded ideas". Storing rejected hypotheses
    lets the agent answer questions like "I used to think X but
    here's why I changed my mind" and prevents re-adopting a
    hypothesis the evidence has already defeated.
    """

    VERIFIED = "verified"
    BELIEVED = "believed"
    UNCERTAIN = "uncertain"
    CONTESTED = "contested"
    SUPERSEDED = "superseded"
    RETRACTED = "retracted"
    REJECTED_HYPOTHESIS = "rejected_hypothesis"


# ── Knowledge Types ──


class KnowledgeType(str, Enum):
    """What kind of knowledge this is. Determines how it should be used."""

    DECLARATIVE = "declarative"
    EPISODIC = "episodic"
    PROCEDURAL = "procedural"
    NORMATIVE = "normative"
    INTENTIONAL = "intentional"


# ── Predicate Definitions ──


class Cardinality(str, Enum):
    ONE = "one"
    MANY = "many"


@dataclass(frozen=True)
class PredicateDef:
    """Formal definition of a predicate."""

    name: str
    domain: frozenset[EntityType]
    range: frozenset[EntityType] | None  # None = literal value
    inverse: str | None = None
    cardinality: Cardinality = Cardinality.MANY
    knowledge_type: KnowledgeType = KnowledgeType.DECLARATIVE
    temporal: bool = False
    implies: tuple[str, ...] = ()
    description: str = ""


@dataclass
class InferenceRule:
    """When a triple with if_predicate is added, create a triple with then_predicate."""

    name: str
    if_predicate: str
    then_predicate: str
    swap: bool = False  # if True, swap subject and object


# ── Ontology State (module-level, loaded once) ──

PREDICATES: dict[str, PredicateDef] = {}
INFERENCE_RULES: list[InferenceRule] = []
_loaded_from: str | None = None  # tracks where ontology was loaded from
_fallback_counts: dict[str, int] = {}  # tracks predicates that fell back to related_to


# ── YAML Loading ──


def _resolve_domain_range(items: list[str] | None) -> frozenset[EntityType] | None:
    """Convert YAML domain/range list to frozenset of EntityType.

    Items can be:
    - Entity type IDs: "person", "organization"
    - Group names: "agent", "artifact", "abstract", "any"
    - null → literal-valued predicate (range only)
    """
    if items is None:
        return None

    types: set[EntityType] = set()
    for item in items:
        item_lower = item.lower()
        # Check if it's a group name
        if item_lower in ENTITY_GROUPS:
            types.update(ENTITY_GROUPS[item_lower])
        else:
            # Try as entity type
            try:
                types.add(EntityType(item_lower))
            except ValueError:
                # Try legacy mapping
                if item_lower in LEGACY_TYPE_MAP:
                    types.add(LEGACY_TYPE_MAP[item_lower])
                else:
                    log.warning("Unknown entity type in ontology YAML: '%s'", item)

    return frozenset(types) if types else frozenset(ALL_TYPES)


def load_from_yaml(yaml_path: Path) -> bool:
    """Load ontology from YAML file. Returns True if successful."""
    global _loaded_from

    try:
        import yaml
    except ImportError:
        log.warning("PyYAML not installed — cannot load ontology YAML")
        return False

    if not yaml_path.exists():
        log.debug("Ontology YAML not found at %s", yaml_path)
        return False

    try:
        data = yaml.safe_load(yaml_path.read_text(encoding="utf-8"))
    except Exception as e:
        log.warning("Failed to parse ontology YAML: %s", e)
        return False

    if not isinstance(data, dict):
        log.warning("Ontology YAML root must be a dict")
        return False

    # Load entity type aliases from YAML into LEGACY_TYPE_MAP
    for etype_def in data.get("entity_types", []):
        eid = etype_def.get("id", "")
        for alias in etype_def.get("aliases", []):
            if alias.lower() not in LEGACY_TYPE_MAP:
                resolved = resolve_entity_type(eid)
                LEGACY_TYPE_MAP[alias.lower()] = resolved

    # Load predicates
    PREDICATES.clear()
    for pred_def in data.get("predicates", []):
        name = pred_def.get("name", "")
        if not name:
            continue

        domain = _resolve_domain_range(pred_def.get("domain"))
        range_ = _resolve_domain_range(pred_def.get("range"))

        kt_str = pred_def.get("knowledge_type", "declarative")
        try:
            kt = KnowledgeType(kt_str)
        except ValueError:
            kt = KnowledgeType.DECLARATIVE

        card_str = pred_def.get("cardinality", "many")
        try:
            card = Cardinality(card_str)
        except ValueError:
            card = Cardinality.MANY

        PREDICATES[name] = PredicateDef(
            name=name,
            domain=domain or frozenset(ALL_TYPES),
            range=range_,
            inverse=pred_def.get("inverse"),
            cardinality=card,
            knowledge_type=kt,
            temporal=bool(pred_def.get("temporal", False)),
            description=pred_def.get("description", ""),
        )

    # Load inference rules
    INFERENCE_RULES.clear()
    for rule_def in data.get("inference_rules", []):
        if_clause = rule_def.get("if", {})
        then_clause = rule_def.get("then", {})
        if not if_clause.get("predicate") or not then_clause.get("predicate"):
            continue
        INFERENCE_RULES.append(InferenceRule(
            name=rule_def.get("name", ""),
            if_predicate=if_clause["predicate"],
            then_predicate=then_clause["predicate"],
            swap=bool(then_clause.get("swap", False)),
        ))

    _loaded_from = str(yaml_path)
    log.debug("Ontology loaded from YAML: %d predicates, %d inference rules",
             len(PREDICATES), len(INFERENCE_RULES))
    return True


def load_ontology(skills_dir: str | Path | None = None) -> None:
    """Load ontology from skills config YAML, falling back to defaults.

    Search order:
    1. skills_dir/_meta/ontology/config/ontology.yaml (user's vault)
    2. defaults/skills/_meta/ontology/config/ontology.yaml (shipped defaults)
    3. Hardcoded fallback (always works, no YAML needed)
    """
    if skills_dir:
        yaml_path = Path(skills_dir) / "_meta" / "ontology" / "config" / "ontology.yaml"
        if load_from_yaml(yaml_path):
            return

    # Try shipped defaults (relative to this file)
    defaults_path = Path(__file__).parent.parent.parent.parent.parent / "defaults" / "skills" / "_meta" / "ontology" / "config" / "ontology.yaml"
    if load_from_yaml(defaults_path):
        return

    # Hardcoded fallback — always works
    _load_hardcoded_defaults()


def _load_hardcoded_defaults() -> None:
    """Hardcoded predicate definitions — fallback when no YAML is available."""
    global _loaded_from

    _ALL = frozenset(ALL_TYPES)
    _AGENTS = frozenset(AGENT_TYPES)
    _PERSONS = frozenset({EntityType.PERSON})
    _ORGS = frozenset({EntityType.ORGANIZATION})
    _ARTIFACTS = frozenset(ARTIFACT_TYPES)
    _TASKS = frozenset({EntityType.TASK})
    _GOALS = frozenset({EntityType.GOAL})
    _PLACES = frozenset({EntityType.PLACE})

    PREDICATES.clear()
    defs = [
        ("works_at", _PERSONS, _ORGS, "employs", Cardinality.ONE),
        ("employs", _ORGS, _PERSONS, "works_at", Cardinality.MANY),
        ("reports_to", _PERSONS, _PERSONS, "manages", Cardinality.ONE),
        ("manages", _PERSONS, _PERSONS, "reports_to", Cardinality.MANY),
        ("knows", _AGENTS, _AGENTS, None, Cardinality.MANY),
        ("collaborates_with", _AGENTS, _AGENTS, None, Cardinality.MANY),
        ("founded", _PERSONS, _ORGS, None, Cardinality.MANY),
        ("member_of", _AGENTS, frozenset({EntityType.ORGANIZATION, EntityType.PROJECT}), None, Cardinality.MANY),
        ("created", _AGENTS, _ARTIFACTS, None, Cardinality.MANY),
        ("owns", _AGENTS, frozenset(ARTIFACT_TYPES | {EntityType.PLACE}), None, Cardinality.MANY),
        ("assigned_to", _TASKS, _AGENTS, "works_on", Cardinality.MANY),
        ("works_on", _AGENTS, _ARTIFACTS, "assigned_to", Cardinality.MANY),
        ("status_is", _ALL, None, None, Cardinality.ONE),
        ("deadline_is", _ARTIFACTS, None, None, Cardinality.ONE),
        ("prefers", _AGENTS, _ALL, None, Cardinality.MANY),
        ("decided", _AGENTS, None, None, Cardinality.MANY),
        ("is_a", _ALL, _ALL, None, Cardinality.MANY),
        ("part_of", _ALL, _ALL, "has_part", Cardinality.MANY),
        ("has_part", _ALL, _ALL, "part_of", Cardinality.MANY),
        ("has_goal", _AGENTS, _GOALS, None, Cardinality.MANY),
        ("interested_in", _AGENTS, frozenset(ABSTRACT_TYPES | ARTIFACT_TYPES), None, Cardinality.MANY),
        ("located_in", _ALL, _PLACES, None, Cardinality.ONE),
        ("uses", _AGENTS, frozenset({EntityType.CONCEPT, EntityType.SKILL, EntityType.DOCUMENT}), None, Cardinality.MANY),
        ("related_to", _ALL, _ALL, None, Cardinality.MANY),
    ]
    for name, domain, range_, inverse, card in defs:
        kt = KnowledgeType.DECLARATIVE
        if name in ("prefers", "decided"):
            kt = KnowledgeType.NORMATIVE
        elif name == "has_goal":
            kt = KnowledgeType.INTENTIONAL
        PREDICATES[name] = PredicateDef(
            name=name, domain=domain, range=range_,
            inverse=inverse, cardinality=card, knowledge_type=kt,
        )

    INFERENCE_RULES.clear()
    _loaded_from = "hardcoded"
    log.debug("Ontology loaded from hardcoded defaults: %d predicates", len(PREDICATES))


# ── Validation ──


@dataclass
class ValidationResult:
    """Result of validating a triple against the ontology."""

    valid: bool
    warnings: list[str] = field(default_factory=list)
    errors: list[str] = field(default_factory=list)
    normalized_predicate: str | None = None
    contradiction: dict | None = None


def validate_triple(
    predicate: str,
    subject_type: EntityType,
    object_type: EntityType | None,
    is_literal: bool = False,
    *,
    strict: bool = False,
) -> ValidationResult:
    """Validate a triple against the ontology schema.

    When *strict* is True, domain/range violations are **errors** that
    set ``valid=False`` — the caller should reject the triple.  When
    False (default, backward-compatible), violations are **warnings**
    and ``valid`` stays True.

    Unknown predicates are always normalized to ``related_to`` regardless
    of strict mode — this is a predicate vocabulary gap, not a type error.
    """
    # Auto-load if not loaded yet
    if not PREDICATES:
        load_ontology()

    result = ValidationResult(valid=True, normalized_predicate=predicate)

    pred_def = PREDICATES.get(predicate)
    was_normalized = False
    if not pred_def:
        # Track this fallback for ontology evolution
        _fallback_counts[predicate] = _fallback_counts.get(predicate, 0) + 1
        result.warnings.append(
            f"Unknown predicate '{predicate}' — mapped to 'related_to'. "
            f"Known predicates: {', '.join(sorted(PREDICATES.keys()))}"
        )
        result.normalized_predicate = "related_to"
        pred_def = PREDICATES.get("related_to")
        was_normalized = True
        if not pred_def:
            return result

    # Skip domain/range checks for normalized predicates — the fallback
    # to related_to already degraded the triple, double-penalizing by
    # also rejecting on the fallback's schema is too aggressive.
    if was_normalized:
        return result

    # Check domain.
    # THING and CONCEPT are treated as "untyped" — they pass domain/range
    # checks because callers often don't know the entity's real type yet.
    _untyped = {EntityType.THING, EntityType.CONCEPT}
    if (
        subject_type not in pred_def.domain
        and subject_type not in _untyped
        and not any(u in pred_def.domain for u in _untyped)
    ):
        domain_ok = any(is_subtype(subject_type, d) for d in pred_def.domain)
        if not domain_ok:
            domain_names = ", ".join(t.value for t in pred_def.domain)
            msg = (
                f"Subject type '{subject_type.value}' not in domain of '{predicate}' "
                f"(expected: {domain_names})"
            )
            if strict:
                result.errors.append(msg)
                result.valid = False
            else:
                result.warnings.append(msg)

    # Check range
    if is_literal:
        # Literals are commonly used as untyped string values (e.g.
        # "has_goal" with "Launch product" as a description rather than
        # a registered goal entity). Always warn, never hard-reject —
        # the value is still useful even if it's not a typed entity.
        if (
            pred_def.range is not None
            and EntityType.THING not in pred_def.range
        ):
            result.warnings.append(
                f"Predicate '{predicate}' expects entity object, got literal"
            )
    elif object_type and pred_def.range is not None:
        if (
            object_type not in pred_def.range
            and object_type not in _untyped
            and not any(u in pred_def.range for u in _untyped)
        ):
            range_ok = any(is_subtype(object_type, r) for r in pred_def.range)
            if not range_ok:
                range_names = ", ".join(t.value for t in pred_def.range)
                msg = (
                    f"Object type '{object_type.value}' not in range of '{predicate}' "
                    f"(expected: {range_names})"
                )
                if strict:
                    result.errors.append(msg)
                    result.valid = False
                else:
                    result.warnings.append(msg)

    return result


# ── Inference ──


def get_inferences(predicate: str) -> list[InferenceRule]:
    """Get all inference rules triggered by a predicate."""
    if not INFERENCE_RULES:
        return []
    return [r for r in INFERENCE_RULES if r.if_predicate == predicate]


# ── Query functions ──


def _ensure_loaded() -> None:
    """Auto-load ontology if not loaded yet."""
    if not PREDICATES:
        load_ontology()


def get_predicate(name: str) -> PredicateDef | None:
    _ensure_loaded()
    return PREDICATES.get(name)


def get_inverse(predicate: str) -> str | None:
    _ensure_loaded()
    pred_def = PREDICATES.get(predicate)
    return pred_def.inverse if pred_def else None


def is_single_cardinality(predicate: str) -> bool:
    _ensure_loaded()
    pred_def = PREDICATES.get(predicate)
    return pred_def.cardinality == Cardinality.ONE if pred_def else False


def get_knowledge_type(predicate: str) -> KnowledgeType:
    _ensure_loaded()
    pred_def = PREDICATES.get(predicate)
    return pred_def.knowledge_type if pred_def else KnowledgeType.DECLARATIVE


def list_predicates() -> list[dict]:
    """List all predicates with definitions (for LLM prompts and UI)."""
    if not PREDICATES:
        load_ontology()
    result = []
    for name, p in sorted(PREDICATES.items()):
        result.append({
            "name": name,
            "domain": [t.value for t in sorted(p.domain, key=lambda x: x.value)],
            "range": [t.value for t in sorted(p.range, key=lambda x: x.value)] if p.range else ["literal"],
            "inverse": p.inverse,
            "cardinality": p.cardinality.value,
            "knowledge_type": p.knowledge_type.value,
            "description": p.description,
        })
    return result


def build_extraction_prompt() -> str:
    """Build the predicate vocabulary section for LLM extraction prompts.

    Generated dynamically from loaded ontology — never hardcoded.
    """
    if not PREDICATES:
        load_ontology()

    lines = []
    for name, p in sorted(PREDICATES.items()):
        if name == "related_to":
            continue  # show last as fallback
        domain = _format_types(p.domain)
        if p.range is None:
            range_ = "literal"
        else:
            range_ = _format_types(p.range)
        inv = f", inverse: {p.inverse}" if p.inverse else ""
        card = ", exclusive" if p.cardinality == Cardinality.ONE else ""
        lines.append(f"- {name} ({domain} → {range_}{inv}{card}): {p.description}")

    lines.append("- related_to (any → any, last resort): Generic relationship")

    # Entity types
    type_lines = []
    for group_name, types in sorted(ENTITY_GROUPS.items()):
        if group_name == "any":
            continue
        type_names = ", ".join(t.value for t in sorted(types, key=lambda x: x.value))
        type_lines.append(f"- {group_name}: {type_names}")

    return (
        "Predicate vocabulary (use these — each has defined domain→range):\n"
        + "\n".join(lines)
        + "\n\nEntity types:\n"
        + "\n".join(type_lines)
        + "\n- literal: for dates, statuses, and other values"
    )


def _format_types(types: frozenset[EntityType]) -> str:
    """Format a set of types concisely for prompt display."""
    if types == frozenset(ALL_TYPES):
        return "any"
    # Check if it matches a group exactly
    for group_name, group_types in ENTITY_GROUPS.items():
        if types == group_types and group_name != "any":
            return group_name
    return "|".join(t.value for t in sorted(types, key=lambda x: x.value))


def get_loaded_from() -> str | None:
    """Return where ontology was loaded from (for diagnostics)."""
    return _loaded_from


def get_fallback_report(min_count: int = 1) -> list[dict]:
    """Get predicates that fell back to related_to, sorted by frequency.

    These are candidates for new predicate definitions.
    """
    return sorted(
        [{"predicate": p, "count": c} for p, c in _fallback_counts.items() if c >= min_count],
        key=lambda x: x["count"],
        reverse=True,
    )


def clear_fallback_counts() -> None:
    """Reset fallback counters (e.g., after adding new predicates)."""
    _fallback_counts.clear()


def reload_ontology(skills_dir: str | Path | None = None) -> dict:
    """Force-reload ontology from YAML. Returns status.

    Called after the agent writes new predicates to ontology.yaml.
    Clears fallback counts for predicates that are now defined.
    """
    old_count = len(PREDICATES)
    load_ontology(skills_dir)
    new_count = len(PREDICATES)

    # Clear fallback counts for predicates that now exist
    resolved = [p for p in list(_fallback_counts.keys()) if p in PREDICATES]
    for p in resolved:
        del _fallback_counts[p]

    return {
        "reloaded": True,
        "source": _loaded_from,
        "predicates": new_count,
        "added": new_count - old_count,
        "resolved_fallbacks": resolved,
    }
