# Copyright 2026 Terrene Foundation
# SPDX-License-Identifier: Apache-2.0

"""
CARE-037: W3C PROV-DM Compatible Provenance Schema for Trust Knowledge Ledger.

Implements the W3C PROV Data Model (PROV-DM) for knowledge provenance tracking.
Provides full lineage tracking for knowledge entries, enabling verification
of knowledge derivation chains and attribution.

W3C PROV-DM defines three core types:
- Entity: Something (physical, digital, conceptual)
- Activity: Something that occurs over a period of time
- Agent: Something that bears responsibility

And relationships between them:
- wasGeneratedBy: Entity was generated by Activity
- wasAttributedTo: Entity was attributed to Agent
- wasAssociatedWith: Activity was associated with Agent
- wasDerivedFrom: Entity was derived from another Entity
- used: Activity used Entity
- wasInformedBy: Activity was informed by another Activity

References:
- W3C PROV-DM: https://www.w3.org/TR/prov-dm/
- W3C PROV-JSON: https://www.w3.org/Submission/prov-json/
"""

from dataclasses import dataclass, field
from datetime import datetime, timezone
from enum import Enum
from typing import Any, Dict, List, Optional, Protocol, Set
from uuid import uuid4


class ProvRelation(Enum):
    """
    W3C PROV-DM relationship types.

    These relationships define how entities, activities, and agents
    are connected in the provenance graph.
    """

    WAS_GENERATED_BY = "wasGeneratedBy"
    WAS_ATTRIBUTED_TO = "wasAttributedTo"
    WAS_ASSOCIATED_WITH = "wasAssociatedWith"
    WAS_DERIVED_FROM = "wasDerivedFrom"
    USED = "used"
    WAS_INFORMED_BY = "wasInformedBy"


class ProvenanceStoreProtocol(Protocol):
    """Protocol for provenance stores used by ProvenanceChain."""

    async def get_provenance(self, entity_id: str) -> Optional["ProvenanceRecord"]:
        """Get provenance record for an entity."""
        ...

    async def knowledge_exists(self, entity_id: str) -> bool:
        """Check if knowledge entity exists."""
        ...


@dataclass
class ProvenanceRecord:
    """
    A provenance record following W3C PROV-DM specification.

    Captures the provenance of a knowledge entry, including:
    - Which activity generated it
    - Which agent is responsible
    - What entities it was derived from
    - Additional attributes

    Attributes:
        record_id: Unique identifier (format: "prov-{uuid_hex[:12]}")
        entity_id: ID of the knowledge entry this provenance is for
        activity_id: ID of the generating activity (format: "activity-{type}-{uuid_hex[:8]}")
        agent_id: ID of the responsible agent
        relations: Map of relation type to list of target IDs
        attributes: Additional PROV attributes
        timestamp: When this provenance record was created (UTC)
    """

    record_id: str
    entity_id: str
    activity_id: str
    agent_id: str
    relations: Dict[str, List[str]] = field(default_factory=dict)
    attributes: Dict[str, Any] = field(default_factory=dict)
    timestamp: datetime = field(default_factory=lambda: datetime.now(timezone.utc))

    @classmethod
    def create_for_knowledge(
        cls,
        knowledge_entry: Any,
        activity: str,
        agent_id: str,
        derived_from: Optional[List[str]] = None,
    ) -> "ProvenanceRecord":
        """
        Create a provenance record for a knowledge entry.

        Automatically sets up PROV-DM relationships:
        - wasGeneratedBy: links entity to activity
        - wasAttributedTo: links entity to agent
        - wasAssociatedWith: links activity to agent
        - wasDerivedFrom: links to source entities (if provided)

        Args:
            knowledge_entry: The KnowledgeEntry this provenance is for
            activity: Type of activity (e.g., "creation", "derivation", "inference")
            agent_id: ID of the responsible agent
            derived_from: Optional list of source entity IDs this was derived from

        Returns:
            ProvenanceRecord with relationships configured per PROV-DM
        """
        record_id = f"prov-{uuid4().hex[:12]}"
        activity_id = f"activity-{activity}-{uuid4().hex[:8]}"
        entity_id = knowledge_entry.entry_id

        # Initialize relations dict
        relations: Dict[str, List[str]] = {}

        # Set up core relationships
        # wasGeneratedBy: entity -> activity
        relations[ProvRelation.WAS_GENERATED_BY.value] = [activity_id]

        # wasAttributedTo: entity -> agent
        relations[ProvRelation.WAS_ATTRIBUTED_TO.value] = [agent_id]

        # wasAssociatedWith: activity -> agent
        relations[ProvRelation.WAS_ASSOCIATED_WITH.value] = [agent_id]

        # wasDerivedFrom: entity -> source entities
        if derived_from:
            relations[ProvRelation.WAS_DERIVED_FROM.value] = list(derived_from)

        # Extract attributes from knowledge entry
        attributes: Dict[str, Any] = {
            "trust_chain_ref": knowledge_entry.trust_chain_ref,
            "content_type": knowledge_entry.content_type.value,
            "confidence_score": knowledge_entry.confidence_score,
        }

        return cls(
            record_id=record_id,
            entity_id=entity_id,
            activity_id=activity_id,
            agent_id=agent_id,
            relations=relations,
            attributes=attributes,
            timestamp=datetime.now(timezone.utc),
        )

    def add_relation(self, relation: ProvRelation, target_id: str) -> None:
        """
        Add a relation to this provenance record.

        Ignores duplicates - if the same relation to the same target
        already exists, it will not be added again.

        Args:
            relation: The type of PROV-DM relation
            target_id: ID of the target entity/activity/agent
        """
        relation_key = relation.value
        if relation_key not in self.relations:
            self.relations[relation_key] = []

        # Only add if not already present (no duplicates)
        if target_id not in self.relations[relation_key]:
            self.relations[relation_key].append(target_id)

    def to_prov_json(self) -> Dict[str, Any]:
        """
        Export to W3C PROV-JSON format.

        Generates a PROV-JSON document following the W3C specification.
        The structure includes separate sections for entities, activities,
        agents, and their relationships.

        Returns:
            Dictionary in PROV-JSON format with sections:
            - entity: Knowledge entry as PROV entity
            - activity: Generating activity
            - agent: Responsible agent
            - wasGeneratedBy: Generation relationships
            - wasAttributedTo: Attribution relationships
            - wasDerivedFrom: Derivation relationships (if any)

        Reference:
            https://www.w3.org/Submission/prov-json/
        """
        prov_json: Dict[str, Any] = {}

        # Entity section
        entity_attrs = {"prov:type": "KnowledgeEntry"}
        entity_attrs.update(self.attributes)
        prov_json["entity"] = {self.entity_id: entity_attrs}

        # Activity section
        prov_json["activity"] = {
            self.activity_id: {
                "prov:type": "KnowledgeActivity",
                "prov:startTime": self.timestamp.isoformat(),
            }
        }

        # Agent section
        prov_json["agent"] = {self.agent_id: {"prov:type": "Agent"}}

        # wasGeneratedBy relationships
        if ProvRelation.WAS_GENERATED_BY.value in self.relations:
            prov_json["wasGeneratedBy"] = {}
            for activity_id in self.relations[ProvRelation.WAS_GENERATED_BY.value]:
                rel_id = f"_:{self.entity_id}_gen"
                prov_json["wasGeneratedBy"][rel_id] = {
                    "prov:entity": self.entity_id,
                    "prov:activity": activity_id,
                }

        # wasAttributedTo relationships
        if ProvRelation.WAS_ATTRIBUTED_TO.value in self.relations:
            prov_json["wasAttributedTo"] = {}
            for agent_id in self.relations[ProvRelation.WAS_ATTRIBUTED_TO.value]:
                rel_id = f"_:{self.entity_id}_attr"
                prov_json["wasAttributedTo"][rel_id] = {
                    "prov:entity": self.entity_id,
                    "prov:agent": agent_id,
                }

        # wasDerivedFrom relationships
        if ProvRelation.WAS_DERIVED_FROM.value in self.relations:
            prov_json["wasDerivedFrom"] = {}
            for i, source_id in enumerate(self.relations[ProvRelation.WAS_DERIVED_FROM.value]):
                rel_id = f"_:{self.entity_id}_der_{i}"
                prov_json["wasDerivedFrom"][rel_id] = {
                    "prov:generatedEntity": self.entity_id,
                    "prov:usedEntity": source_id,
                }

        return prov_json

    def to_dict(self) -> Dict[str, Any]:
        """
        Serialize the provenance record to a dictionary.

        Returns:
            Dictionary with all provenance record fields
        """
        return {
            "record_id": self.record_id,
            "entity_id": self.entity_id,
            "activity_id": self.activity_id,
            "agent_id": self.agent_id,
            "relations": {k: list(v) for k, v in self.relations.items()},
            "attributes": dict(self.attributes),
            "timestamp": self.timestamp.isoformat(),
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "ProvenanceRecord":
        """
        Deserialize a provenance record from a dictionary.

        Args:
            data: Dictionary with provenance record fields

        Returns:
            ProvenanceRecord instance
        """
        # Parse timestamp
        timestamp = data.get("timestamp")
        if isinstance(timestamp, str):
            timestamp = datetime.fromisoformat(timestamp)
        elif timestamp is None:
            timestamp = datetime.now(timezone.utc)

        return cls(
            record_id=data["record_id"],
            entity_id=data["entity_id"],
            activity_id=data["activity_id"],
            agent_id=data["agent_id"],
            relations={k: list(v) for k, v in data.get("relations", {}).items()},
            attributes=dict(data.get("attributes", {})),
            timestamp=timestamp,
        )


class ProvenanceChain:
    """
    Traverses and verifies provenance chains.

    Provides methods to follow wasDerivedFrom relationships to build
    complete lineage trees and verify the integrity of derivation chains.
    """

    def __init__(self, provenance_store: ProvenanceStoreProtocol):
        """
        Initialize with a provenance store.

        Args:
            provenance_store: Object implementing get_provenance() and
                            knowledge_exists() async methods
        """
        self._store = provenance_store

    async def get_lineage(self, entity_id: str, max_depth: int = 10) -> List[ProvenanceRecord]:
        """
        Get the complete lineage of an entity using BFS traversal.

        Follows wasDerivedFrom relationships backwards to find all
        source entities in the derivation chain.

        Args:
            entity_id: ID of the entity to trace lineage for
            max_depth: Maximum depth to traverse (default 10)

        Returns:
            List of ProvenanceRecord objects in BFS order,
            starting with the given entity
        """
        lineage: List[ProvenanceRecord] = []
        visited: Set[str] = set()
        queue: List[tuple[str, int]] = [(entity_id, 0)]

        while queue:
            current_id, depth = queue.pop(0)

            if current_id in visited or depth > max_depth:
                continue

            visited.add(current_id)

            # Get provenance for current entity
            record = await self._store.get_provenance(current_id)
            if record is None:
                continue

            lineage.append(record)

            # Add derived-from sources to queue
            derived_from = record.relations.get(ProvRelation.WAS_DERIVED_FROM.value, [])
            for source_id in derived_from:
                if source_id not in visited:
                    queue.append((source_id, depth + 1))

        return lineage

    async def verify_chain_integrity(self, entity_id: str, max_depth: int = 100) -> bool:
        """
        Verify that all entities in the derivation chain exist.

        Traverses the complete lineage and checks that every
        referenced entity exists in the knowledge store.

        CARE-044: max_depth bounds traversal to prevent DoS via
        deeply nested provenance graphs.

        Args:
            entity_id: ID of the entity to verify chain for
            max_depth: Maximum traversal depth (default 100)

        Returns:
            True if all entities in the chain exist, False otherwise
        """
        visited: Set[str] = set()
        queue: List[tuple[str, int]] = [(entity_id, 0)]

        while queue:
            current_id, depth = queue.pop(0)

            if current_id in visited or depth > max_depth:
                continue

            visited.add(current_id)

            # Check if entity exists
            exists = await self._store.knowledge_exists(current_id)
            if not exists:
                return False

            # Get provenance to find derived-from sources
            record = await self._store.get_provenance(current_id)
            if record is None:
                # If there's no provenance record but entity exists, that's ok
                continue

            # Add derived-from sources to queue
            derived_from = record.relations.get(ProvRelation.WAS_DERIVED_FROM.value, [])
            for source_id in derived_from:
                if source_id not in visited:
                    queue.append((source_id, depth + 1))

        return True


class InMemoryProvenanceStore:
    """
    In-memory provenance store for testing and development.

    Stores provenance records in memory with support for querying
    by derived-from relationships.
    """

    def __init__(self) -> None:
        """Initialize empty store."""
        self._records: Dict[str, ProvenanceRecord] = {}

    async def store(self, record: ProvenanceRecord) -> None:
        """
        Store a provenance record.

        Args:
            record: ProvenanceRecord to store
        """
        self._records[record.entity_id] = record

    async def get_provenance(self, entity_id: str) -> Optional[ProvenanceRecord]:
        """
        Get provenance record for an entity.

        Args:
            entity_id: ID of the entity

        Returns:
            ProvenanceRecord if found, None otherwise
        """
        return self._records.get(entity_id)

    async def query(self, derived_from: Optional[str] = None) -> List[ProvenanceRecord]:
        """
        Query provenance records.

        Args:
            derived_from: If provided, find records that have wasDerivedFrom
                         referencing this entity

        Returns:
            List of matching ProvenanceRecord objects
        """
        if derived_from is None:
            return list(self._records.values())

        # Find records that derive from the given entity
        results: List[ProvenanceRecord] = []
        for record in self._records.values():
            sources = record.relations.get(ProvRelation.WAS_DERIVED_FROM.value, [])
            if derived_from in sources:
                results.append(record)

        return results

    async def knowledge_exists(self, entity_id: str) -> bool:
        """
        Check if a knowledge entity exists (has provenance record).

        Args:
            entity_id: ID of the entity to check

        Returns:
            True if entity has a provenance record, False otherwise
        """
        return entity_id in self._records
