"""
InsAIts SDK - AI Lineage & Anomaly Tracing Oracle
==================================================
Full message provenance tracking for compliance (GDPR, HIPAA, SEC).
SHA256 hashing for tamper detection, parent-child relationships,
and lineage drift detection using embeddings.

Key Features:
- Complete audit trail with cryptographic hashes
- Parent-child message chain tracking
- Lineage drift detection (semantic divergence)
- Compliance exports (JSON, CSV)
- Integration with hallucination detection

Usage:
    from insa_its.lineage_oracle import LineageOracle, LineageRecord

    oracle = LineageOracle()
    record = oracle.record_message(text, sender, receiver, llm_id)

    # Verify integrity
    is_valid = oracle.verify_chain_integrity()

    # Export for compliance
    oracle.export_audit_trail("audit.json")
"""

import hashlib
import json
import csv
import time
import logging
from dataclasses import dataclass, field, asdict
from typing import Dict, List, Optional, Any, Set, Tuple
from datetime import datetime
from enum import Enum
import numpy as np

logger = logging.getLogger(__name__)


class LineageStatus(Enum):
    """Status of a lineage record."""
    VALID = "valid"
    TAMPERED = "tampered"
    ORPHANED = "orphaned"  # Parent not found
    ANOMALOUS = "anomalous"  # Contains anomaly


@dataclass
class LineageRecord:
    """
    Immutable record of a message in the lineage chain.

    All fields are hashed together to create a tamper-evident record.
    """
    message_id: str
    parent_id: Optional[str]  # Previous message in chain
    text_hash: str  # SHA256 of message text (privacy: don't store text)
    timestamp: float  # Unix timestamp with microseconds
    agent_id: str
    llm_id: str
    sender: str
    receiver: str
    input_refs: List[str] = field(default_factory=list)  # Source message IDs
    embedding_hash: Optional[str] = None  # Hash of embedding for drift detection
    anomaly_types: List[str] = field(default_factory=list)  # Detected anomalies
    metadata: Dict[str, Any] = field(default_factory=dict)
    record_hash: str = ""  # Hash of this entire record

    def __post_init__(self):
        """Calculate record hash after initialization."""
        if not self.record_hash:
            self.record_hash = self._calculate_hash()

    def _calculate_hash(self) -> str:
        """Calculate SHA256 hash of all record fields."""
        data = {
            "message_id": self.message_id,
            "parent_id": self.parent_id,
            "text_hash": self.text_hash,
            "timestamp": self.timestamp,
            "agent_id": self.agent_id,
            "llm_id": self.llm_id,
            "sender": self.sender,
            "receiver": self.receiver,
            "input_refs": sorted(self.input_refs),
            "embedding_hash": self.embedding_hash,
            "anomaly_types": sorted(self.anomaly_types),
        }
        serialized = json.dumps(data, sort_keys=True)
        return hashlib.sha256(serialized.encode()).hexdigest()

    def verify(self) -> bool:
        """Verify record hasn't been tampered with."""
        expected_hash = self._calculate_hash()
        return self.record_hash == expected_hash

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for serialization."""
        return asdict(self)

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "LineageRecord":
        """Create record from dictionary."""
        return cls(**data)


@dataclass
class LineageAnomaly:
    """An anomaly detected in the lineage chain."""
    anomaly_id: str
    source_message_id: str  # Where anomaly originated
    propagation_path: List[str]  # Message IDs it propagated through
    anomaly_type: str
    severity: str
    first_detected: float
    description: str


class LineageOracle:
    """
    AI Lineage & Anomaly Tracing Oracle.

    Tracks complete message provenance for compliance and debugging.
    Provides tamper detection, chain integrity verification, and
    compliance-ready exports.
    """

    def __init__(
        self,
        enable_drift_detection: bool = True,
        drift_threshold: float = 0.65,
        max_records: int = 10000,
    ):
        self.records: Dict[str, LineageRecord] = {}
        self.chains: Dict[str, List[str]] = {}  # agent_id -> [message_ids]
        self.anomalies: Dict[str, LineageAnomaly] = {}
        self.enable_drift_detection = enable_drift_detection
        self.drift_threshold = drift_threshold
        self.max_records = max_records

        # For embedding-based drift detection
        self._embedding_cache: Dict[str, np.ndarray] = {}

    def record_message(
        self,
        text: str,
        sender: str,
        receiver: str,
        llm_id: str,
        message_id: Optional[str] = None,
        parent_id: Optional[str] = None,
        input_refs: Optional[List[str]] = None,
        embedding: Optional[np.ndarray] = None,
        anomaly_types: Optional[List[str]] = None,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> LineageRecord:
        """
        Record a message in the lineage chain.

        Args:
            text: Message text (will be hashed, not stored)
            sender: Sender agent ID
            receiver: Receiver agent ID
            llm_id: LLM model identifier
            message_id: Optional custom message ID
            parent_id: Previous message in conversation
            input_refs: Source message IDs this derives from
            embedding: Message embedding for drift detection
            anomaly_types: Any detected anomalies
            metadata: Additional metadata

        Returns:
            Created LineageRecord
        """
        import uuid

        # Generate message ID if not provided
        if message_id is None:
            message_id = str(uuid.uuid4())

        # Hash the text (privacy: don't store raw text)
        text_hash = hashlib.sha256(text.encode()).hexdigest()

        # Hash embedding if provided
        embedding_hash = None
        if embedding is not None:
            embedding_bytes = embedding.tobytes()
            embedding_hash = hashlib.sha256(embedding_bytes).hexdigest()[:16]
            self._embedding_cache[message_id] = embedding

        # Auto-detect parent if not specified
        if parent_id is None:
            agent_key = f"{sender}:{receiver}"
            if agent_key in self.chains and self.chains[agent_key]:
                parent_id = self.chains[agent_key][-1]

        # Create record
        record = LineageRecord(
            message_id=message_id,
            parent_id=parent_id,
            text_hash=text_hash,
            timestamp=time.time(),
            agent_id=sender,
            llm_id=llm_id,
            sender=sender,
            receiver=receiver,
            input_refs=input_refs or [],
            embedding_hash=embedding_hash,
            anomaly_types=anomaly_types or [],
            metadata=metadata or {},
        )

        # Store record
        self.records[message_id] = record

        # Update chain
        agent_key = f"{sender}:{receiver}"
        if agent_key not in self.chains:
            self.chains[agent_key] = []
        self.chains[agent_key].append(message_id)

        # Check for drift if enabled
        if self.enable_drift_detection and parent_id and embedding is not None:
            self._check_drift(record, embedding)

        # Cleanup old records if over limit
        if len(self.records) > self.max_records:
            self._cleanup_old_records()

        logger.debug(f"Recorded message {message_id} in lineage")
        return record

    def _check_drift(self, record: LineageRecord, embedding: np.ndarray):
        """Check for semantic drift from parent message."""
        if record.parent_id not in self._embedding_cache:
            return

        parent_embedding = self._embedding_cache[record.parent_id]

        # Calculate cosine similarity
        similarity = np.dot(embedding, parent_embedding) / (
            np.linalg.norm(embedding) * np.linalg.norm(parent_embedding)
        )

        if similarity < self.drift_threshold:
            logger.warning(
                f"Lineage drift detected: {record.message_id} "
                f"(similarity={similarity:.2f} < {self.drift_threshold})"
            )
            # Add to anomalies
            if "LINEAGE_DRIFT" not in record.anomaly_types:
                record.anomaly_types.append("LINEAGE_DRIFT")
                record.metadata["drift_similarity"] = float(similarity)

    def _cleanup_old_records(self):
        """Remove oldest records to stay under limit."""
        if len(self.records) <= self.max_records:
            return

        # Sort by timestamp
        sorted_ids = sorted(
            self.records.keys(),
            key=lambda x: self.records[x].timestamp
        )

        # Remove oldest 10%
        to_remove = sorted_ids[:len(sorted_ids) // 10]
        for msg_id in to_remove:
            del self.records[msg_id]
            if msg_id in self._embedding_cache:
                del self._embedding_cache[msg_id]

    def get_record(self, message_id: str) -> Optional[LineageRecord]:
        """Get a specific lineage record."""
        return self.records.get(message_id)

    def get_data_provenance(
        self,
        message_id: str,
        max_depth: int = 50,
    ) -> Dict[str, Any]:
        """
        Trace a message back to its origin.

        Args:
            message_id: Message to trace
            max_depth: Maximum chain depth to traverse

        Returns:
            Dict with full provenance chain and metadata
        """
        if message_id not in self.records:
            return {"error": "Message not found", "message_id": message_id}

        chain = []
        current_id = message_id
        visited = set()
        depth = 0

        while current_id and depth < max_depth:
            if current_id in visited:
                break  # Cycle detected
            visited.add(current_id)

            record = self.records.get(current_id)
            if not record:
                break

            chain.append({
                "message_id": record.message_id,
                "agent": record.agent_id,
                "llm": record.llm_id,
                "timestamp": datetime.fromtimestamp(record.timestamp).isoformat(),
                "has_anomalies": len(record.anomaly_types) > 0,
                "anomalies": record.anomaly_types,
            })

            current_id = record.parent_id
            depth += 1

        # Also trace input_refs (lateral sources)
        lateral_sources = []
        target_record = self.records.get(message_id)
        if target_record:
            for ref_id in target_record.input_refs:
                if ref_id in self.records:
                    ref_record = self.records[ref_id]
                    lateral_sources.append({
                        "message_id": ref_record.message_id,
                        "agent": ref_record.agent_id,
                        "llm": ref_record.llm_id,
                    })

        return {
            "message_id": message_id,
            "chain_length": len(chain),
            "chain": chain,
            "lateral_sources": lateral_sources,
            "origin": chain[-1] if chain else None,
        }

    def verify_chain_integrity(
        self,
        start_message_id: Optional[str] = None,
    ) -> Dict[str, Any]:
        """
        Verify integrity of the lineage chain.

        Args:
            start_message_id: Start from specific message (or verify all)

        Returns:
            Verification results
        """
        if start_message_id:
            records_to_check = [self.records.get(start_message_id)]
            if records_to_check[0] is None:
                return {"valid": False, "error": "Message not found"}
        else:
            records_to_check = list(self.records.values())

        results = {
            "valid": True,
            "total_records": len(records_to_check),
            "valid_count": 0,
            "tampered_count": 0,
            "orphaned_count": 0,
            "issues": [],
        }

        for record in records_to_check:
            # Check record hash
            if not record.verify():
                results["valid"] = False
                results["tampered_count"] += 1
                results["issues"].append({
                    "message_id": record.message_id,
                    "issue": "TAMPERED",
                    "description": "Record hash mismatch",
                })
                continue

            # Check parent exists (if specified)
            if record.parent_id and record.parent_id not in self.records:
                results["orphaned_count"] += 1
                results["issues"].append({
                    "message_id": record.message_id,
                    "issue": "ORPHANED",
                    "description": f"Parent {record.parent_id} not found",
                })
                # Orphaned is a warning, not invalid
                results["valid_count"] += 1
                continue

            results["valid_count"] += 1

        return results

    def detect_lineage_anomalies(self) -> List[LineageAnomaly]:
        """
        Detect anomalies in lineage chains.

        Returns:
            List of detected LineageAnomaly objects
        """
        anomalies = []

        # Track anomaly propagation
        anomaly_sources: Dict[str, Set[str]] = {}  # anomaly_type -> source_message_ids

        for msg_id, record in self.records.items():
            if not record.anomaly_types:
                continue

            for anomaly_type in record.anomaly_types:
                # Track this as a potential source
                if anomaly_type not in anomaly_sources:
                    anomaly_sources[anomaly_type] = set()
                anomaly_sources[anomaly_type].add(msg_id)

        # Trace propagation paths
        for anomaly_type, source_ids in anomaly_sources.items():
            for source_id in source_ids:
                # Find messages that derive from this source
                propagation = self._trace_propagation(source_id)

                if propagation:
                    anomaly = LineageAnomaly(
                        anomaly_id=f"{anomaly_type}:{source_id[:8]}",
                        source_message_id=source_id,
                        propagation_path=propagation,
                        anomaly_type=anomaly_type,
                        severity=self._get_severity(anomaly_type),
                        first_detected=self.records[source_id].timestamp,
                        description=f"{anomaly_type} originated at {source_id[:8]} and propagated to {len(propagation)} messages",
                    )
                    anomalies.append(anomaly)
                    self.anomalies[anomaly.anomaly_id] = anomaly

        return anomalies

    def _trace_propagation(self, source_id: str) -> List[str]:
        """Trace how a message propagated through the system."""
        propagation = []

        for msg_id, record in self.records.items():
            if msg_id == source_id:
                continue
            # Check if this message derives from source
            if record.parent_id == source_id or source_id in record.input_refs:
                propagation.append(msg_id)

        return propagation

    def _get_severity(self, anomaly_type: str) -> str:
        """Get severity level for anomaly type."""
        high_severity = {
            "FACT_CONTRADICTION", "PHANTOM_CITATION", "TAMPERED",
            "HALLUCINATION", "LINEAGE_DRIFT"
        }
        if anomaly_type in high_severity:
            return "high"
        return "medium"

    def export_audit_trail(
        self,
        filepath: str,
        format: str = "json",
        include_metadata: bool = True,
    ) -> Dict[str, Any]:
        """
        Export lineage data for compliance/audit purposes.

        Args:
            filepath: Output file path
            format: "json" or "csv"
            include_metadata: Include metadata in export

        Returns:
            Export summary
        """
        records_data = []

        for record in self.records.values():
            data = {
                "message_id": record.message_id,
                "parent_id": record.parent_id,
                "text_hash": record.text_hash,
                "timestamp": datetime.fromtimestamp(record.timestamp).isoformat(),
                "timestamp_unix": record.timestamp,
                "agent_id": record.agent_id,
                "llm_id": record.llm_id,
                "sender": record.sender,
                "receiver": record.receiver,
                "input_refs": ",".join(record.input_refs),
                "anomaly_types": ",".join(record.anomaly_types),
                "record_hash": record.record_hash,
                "is_valid": record.verify(),
            }

            if include_metadata and record.metadata:
                data["metadata"] = json.dumps(record.metadata)

            records_data.append(data)

        # Sort by timestamp
        records_data.sort(key=lambda x: x["timestamp_unix"])

        if format.lower() == "json":
            export_data = {
                "export_timestamp": datetime.now().isoformat(),
                "total_records": len(records_data),
                "integrity_check": self.verify_chain_integrity(),
                "records": records_data,
            }

            with open(filepath, "w") as f:
                json.dump(export_data, f, indent=2)

        elif format.lower() == "csv":
            if records_data:
                fieldnames = list(records_data[0].keys())
                with open(filepath, "w", newline="") as f:
                    writer = csv.DictWriter(f, fieldnames=fieldnames)
                    writer.writeheader()
                    writer.writerows(records_data)

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

        return {
            "filepath": filepath,
            "format": format,
            "records_exported": len(records_data),
            "timestamp": datetime.now().isoformat(),
        }

    def get_summary(self) -> Dict[str, Any]:
        """Get summary of lineage tracking."""
        total_records = len(self.records)
        anomalous_count = sum(
            1 for r in self.records.values() if r.anomaly_types
        )

        # Count by agent
        by_agent: Dict[str, int] = {}
        for record in self.records.values():
            by_agent[record.agent_id] = by_agent.get(record.agent_id, 0) + 1

        # Count by LLM
        by_llm: Dict[str, int] = {}
        for record in self.records.values():
            by_llm[record.llm_id] = by_llm.get(record.llm_id, 0) + 1

        return {
            "total_records": total_records,
            "anomalous_records": anomalous_count,
            "anomaly_rate": f"{anomalous_count / total_records:.1%}" if total_records > 0 else "0%",
            "unique_agents": len(by_agent),
            "by_agent": by_agent,
            "unique_llms": len(by_llm),
            "by_llm": by_llm,
            "chain_count": len(self.chains),
            "drift_detection_enabled": self.enable_drift_detection,
        }

    def clear(self):
        """Clear all lineage records."""
        self.records.clear()
        self.chains.clear()
        self.anomalies.clear()
        self._embedding_cache.clear()
        logger.info("Lineage oracle cleared")
