import hashlib
import json
import uuid
from datetime import datetime
from functools import wraps
from typing import Any, Callable, Optional

from .audit_log import AuditLog
from .models import AuditEvent, ContextRequest, Decision, SensitivityLevel
from .policy_engine import PolicyEngine


class ContextGuard:
    """
    AutoPIL ContextGuard — the core enforcement point.

    Wraps any retrieval function and enforces:
      - Policy-based access control at retrieval time
      - Immutable audit logging of every context event
      - Context hash for provenance
    """

    def __init__(
        self,
        policy_path: str,
        audit_db: str = "autopil_audit.db",
        tenant_id: str = "default",
    ):
        self.policy_engine = PolicyEngine(policy_path)
        self.audit_log = AuditLog(audit_db)
        self.tenant_id = tenant_id

    def protect(
        self,
        agent_role: str,
        user_id: str,
        source_id: str,
        sensitivity_level: SensitivityLevel = SensitivityLevel.MEDIUM,
        session_id: Optional[str] = None,
    ):
        """
        Decorator that intercepts context retrieval calls and enforces governance.

        Usage:
            @guard.protect(agent_role="loan_underwriter", user_id="u1", source_id="account_summaries")
            def retrieve(query: str) -> list:
                return vectorstore.search(query)
        """

        def decorator(func: Callable) -> Callable:
            @wraps(func)
            def wrapper(*args, **kwargs):
                query = args[0] if args else kwargs.get("query", "")
                sid = session_id or str(uuid.uuid4())

                request = ContextRequest(
                    query=str(query),
                    agent_role=agent_role,
                    user_id=user_id,
                    source_id=source_id,
                    sensitivity_level=sensitivity_level,
                    session_id=sid,
                )

                # Cross-agent isolation — session belongs to first agent that used it
                owner = self.audit_log.get_session_owner(sid, self.tenant_id)
                if owner and owner != agent_role:
                    decision   = Decision.DENY
                    policy_name = "cross_agent_isolation"
                    reason      = (
                        f"Session '{sid[:8]}…' is owned by '{owner}' — "
                        f"'{agent_role}' cannot access another agent's context"
                    )
                else:
                    decision, policy_name, reason = self.policy_engine.evaluate(request)

                context = None
                context_hash = None

                if decision == Decision.ALLOW:
                    context = func(*args, **kwargs)
                    context_hash = hashlib.sha256(
                        json.dumps(str(context), sort_keys=True).encode()
                    ).hexdigest()[:16]

                self.audit_log.record(
                    AuditEvent(
                        event_id=str(uuid.uuid4()),
                        session_id=sid,
                        agent_role=agent_role,
                        user_id=user_id,
                        source_id=source_id,
                        query=str(query),
                        decision=decision,
                        policy_name=policy_name,
                        reason=reason,
                        timestamp=datetime.utcnow(),
                        context_hash=context_hash,
                    ),
                    tenant_id=self.tenant_id,
                )

                if decision == Decision.DENY:
                    raise PermissionError(
                        f"[AutoPIL] DENIED | source={source_id!r} | "
                        f"agent={agent_role!r} | {reason}"
                    )

                return context

            return wrapper

        return decorator

    def get_audit_trail(self, session_id: str) -> list[AuditEvent]:
        return self.audit_log.get_session_events(session_id, self.tenant_id)
