import contextvars
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 .db.sqlite import SQLiteLineageStore
from .models import AgentAction, AuditEvent, ContextRequest, Decision, SensitivityLevel
from .policy_engine import PolicyEngine
from . import telemetry

# Optional — only present when autopil[catalog-*] extras are installed
try:
    from .catalog.enricher import ClassificationEnricher as _ClassificationEnricher
except ImportError:
    _ClassificationEnricher = None  # type: ignore[assignment,misc]


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

    Both sync and async retrieval functions are supported:

        # Sync
        @guard.protect(agent_role="analyst", user_id="u1", source_id="reports")
        def fetch(query: str) -> list:
            return vectorstore.search(query)

        # Async
        @guard.protect_async(agent_role="analyst", user_id="u1", source_id="reports")
        async def fetch(query: str) -> list:
            return await vectorstore.asearch(query)
    """

    _source_type: str = "sdk"  # overridden by integration-specific subclasses

    def __init__(
        self,
        policy_path: str,
        audit_db: str = "autopil_audit.db",
        tenant_id: str = "default",
        catalog_enricher=None,
    ):
        self.policy_engine = PolicyEngine(policy_path)
        self.audit_log = AuditLog(audit_db)
        self.lineage_store = SQLiteLineageStore(audit_db)
        self.tenant_id = tenant_id
        self._catalog_enricher = catalog_enricher  # ClassificationEnricher | None
        # ContextVar is safe across both threads and concurrent async Tasks.
        # Each asyncio Task gets its own context copy, so two concurrent Tasks
        # using the same guard instance never overwrite each other's last_event_id.
        self._last_event_id: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar(
            f"autopil_last_event_id_{id(self)}", default=None
        )

    @property
    def last_event_id(self) -> Optional[str]:
        """event_id of the most recent evaluation on this thread or coroutine."""
        return self._last_event_id.get()

    # ── shared internal helpers ──────────────────────────────────────────────

    def _evaluate_request(
        self,
        agent_role: str,
        user_id: str,
        source_id: str,
        sensitivity_level: SensitivityLevel,
        sid: str,
        query: str,
    ) -> tuple[Decision, str, str]:
        """
        Run cross-agent isolation check + policy evaluation.
        Returns (decision, policy_name, reason).
        Called by both protect() and protect_async().
        """
        request = ContextRequest(
            query=query,
            agent_role=agent_role,
            user_id=user_id,
            source_id=source_id,
            sensitivity_level=sensitivity_level,
            session_id=sid,
        )

        # Catalog enrichment — raises sensitivity if catalog says higher (fail-open)
        if self._catalog_enricher is not None:
            request = self._catalog_enricher.enrich(request, self.tenant_id)

        with telemetry.evaluate_span(
            agent_role=agent_role,
            source_id=source_id,
            tenant_id=self.tenant_id,
            session_id=sid,
            sensitivity_level=sensitivity_level.value,
        ) as otel_span:
            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)
            otel_span.set_result(decision=decision.value, policy_name=policy_name, reason=reason)

        return decision, policy_name, reason

    def _record_event(
        self,
        *,
        agent_role: str,
        user_id: str,
        source_id: str,
        query: str,
        decision: Decision,
        policy_name: str,
        reason: str,
        sensitivity_level: SensitivityLevel,
        sid: str,
        context_hash: Optional[str],
        event_id: Optional[str] = None,
        source_type: Optional[str] = None,
    ) -> str:
        """
        Write the audit event and update last_event_id for this context.
        Returns the new event_id.
        Called by both protect() and protect_async().

        event_id: pre-generated ID (e.g. from middleware); auto-generated if omitted.
        source_type: override for self._source_type (e.g. "api" from middleware).
        """
        event_id = event_id or str(uuid.uuid4())
        self._last_event_id.set(event_id)
        event = AuditEvent(
            event_id=event_id,
            session_id=sid,
            agent_role=agent_role,
            user_id=user_id,
            source_id=source_id,
            query=query,
            decision=decision,
            policy_name=policy_name,
            reason=reason,
            timestamp=datetime.utcnow(),
            context_hash=context_hash,
            sensitivity_level=sensitivity_level.value,
            source_type=source_type if source_type is not None else self._source_type,
        )
        self.audit_log.record(event, tenant_id=self.tenant_id)

        # Catalog write-back — fire-and-forget daemon thread (fail-open)
        if self._catalog_enricher is not None:
            self._catalog_enricher.write_back(event, self.tenant_id)

        return event_id

    # ── public decorators ────────────────────────────────────────────────────

    def protect(
        self,
        agent_role: str,
        user_id: str,
        source_id: str,
        sensitivity_level: SensitivityLevel = SensitivityLevel.MEDIUM,
        session_id: Optional[str] = None,
    ):
        """
        Decorator for sync retrieval functions.

        Usage:
            @guard.protect(agent_role="analyst", user_id="u1", source_id="reports")
            def fetch(query: str) -> list:
                return vectorstore.search(query)
        """
        def decorator(func: Callable) -> Callable:
            @wraps(func)
            def wrapper(*args, **kwargs):
                query = str(args[0]) if args else str(kwargs.get("query", ""))
                sid   = session_id or str(uuid.uuid4())

                decision, policy_name, reason = self._evaluate_request(
                    agent_role, user_id, source_id, sensitivity_level, sid, query
                )

                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._record_event(
                    agent_role=agent_role, user_id=user_id, source_id=source_id,
                    query=query, decision=decision, policy_name=policy_name,
                    reason=reason, sensitivity_level=sensitivity_level,
                    sid=sid, context_hash=context_hash,
                )

                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 protect_async(
        self,
        agent_role: str,
        user_id: str,
        source_id: str,
        sensitivity_level: SensitivityLevel = SensitivityLevel.MEDIUM,
        session_id: Optional[str] = None,
    ):
        """
        Decorator for async retrieval functions.

        Policy evaluation and audit logging are performed synchronously
        (sub-millisecond for in-memory policy + SQLite). The protected
        function itself is awaited, so truly async retrievals are fully
        supported without blocking the event loop.

        For workloads where the audit write latency matters, use the
        Postgres backend (non-blocking at scale) or wrap audit_log.record
        in asyncio.to_thread() in your own subclass.

        Usage:
            @guard.protect_async(agent_role="analyst", user_id="u1", source_id="reports")
            async def fetch(query: str) -> list:
                return await vectorstore.asearch(query)

        Concurrent safety:
            last_event_id is stored in a ContextVar. Each asyncio Task gets
            its own context copy, so two concurrent Tasks using the same guard
            instance never overwrite each other's last_event_id.
        """
        def decorator(func: Callable) -> Callable:
            @wraps(func)
            async def wrapper(*args, **kwargs):
                query = str(args[0]) if args else str(kwargs.get("query", ""))
                sid   = session_id or str(uuid.uuid4())

                decision, policy_name, reason = self._evaluate_request(
                    agent_role, user_id, source_id, sensitivity_level, sid, query
                )

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

                self._record_event(
                    agent_role=agent_role, user_id=user_id, source_id=source_id,
                    query=query, decision=decision, policy_name=policy_name,
                    reason=reason, sensitivity_level=sensitivity_level,
                    sid=sid, context_hash=context_hash,
                )

                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

    # ── lineage ──────────────────────────────────────────────────────────────

    def record_action(
        self,
        event_id: str,
        action_type: str,
        detail: dict,
        outcome: Optional[str] = None,
    ) -> AgentAction:
        """
        Record an agent action driven by a specific audit event.

        Call this after a successful @guard.protect() or @guard.protect_async()
        invocation, passing guard.last_event_id as the event_id.

        Example:
            result = await retrieve_credit_score("cust_abc")
            guard.record_action(
                event_id=guard.last_event_id,
                action_type="loan_decision",
                detail={"outcome": "approved", "amount": 250000},
            )
        """
        events = self.audit_log.get_all_events(self.tenant_id)
        event  = next((e for e in events if e.event_id == event_id), None)
        if event is None:
            raise ValueError(f"No audit event found with event_id='{event_id}'")
        action = AgentAction(
            action_id=str(uuid.uuid4()),
            event_id=event_id,
            session_id=event.session_id,
            agent_role=event.agent_role,
            action_type=action_type,
            action_detail=detail,
            outcome=outcome,
            timestamp=datetime.utcnow(),
        )
        self.lineage_store.record_action(action, self.tenant_id)
        return action

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