"""
Agendex Client - Core SDK for AI Agent Governance

The AgendexClient handles both shadow and enforce modes:
- Shadow mode: All actions allowed, full activity captured for analysis
- Enforce mode: Actions evaluated against policy, blocked/approved accordingly
"""
from __future__ import annotations

import logging
import os
from time import perf_counter
from typing import Any, Callable, Dict, List, Optional

import requests

from .errors import (
    AgendexError,
    ConfigurationError,
    ConnectionError,
    DeniedError,
    PendingApprovalError,
    ShadowDecisionError,
)
from .activity import build_activity_record, extract_resources

logger = logging.getLogger("agendex")


class AgendexClient:
    """
    Client for integrating AI agents with Agendex governance proxy.
    
    Handles the full governance lifecycle:
    1. Activity capture (shadow mode) or policy evaluation (enforce mode)
    2. Token-scoped action invocation
    3. Approval workflow integration
    
    Configuration via environment variables:
        AGENDEX_URL: Proxy URL (default: http://localhost:8000)
        AGENDEX_AGENT_ID: Agent identifier (required)
        AGENDEX_TOKEN: Bearer token for authentication (required)
        AGENDEX_TENANT_ID: Optional tenant identifier
    
    Example:
        client = AgendexClient()
        
        try:
            result = client.invoke(
                action="tool.fetch_revenue",
                params={"month": "2025-01"},
                task="monthly_report"
            )
        except DeniedError as e:
            print(f"Blocked: {e.reason}")
        except PendingApprovalError as e:
            print(f"Needs approval: {e.approval_id}")
    """
    
    def __init__(
        self,
        base_url: Optional[str] = None,
        agent_id: Optional[str] = None,
        token: Optional[str] = None,
        tenant_id: Optional[str] = None,
        mode: str = "auto",
        timeout: int = 120,
        session: Optional[requests.Session] = None,
    ):
        """
        Initialize the Agendex client.
        
        Args:
            base_url: Agendex proxy URL (default: AGENDEX_URL env or http://localhost:8000)
            agent_id: Agent identifier (default: AGENDEX_AGENT_ID env)
            token: Bearer token for auth (default: AGENDEX_TOKEN env)
            tenant_id: Optional tenant identifier (default: AGENDEX_TENANT_ID env)
            mode: Operating mode - "shadow", "enforce", or "auto" (fetch from server)
            timeout: Request timeout in seconds
            session: Optional requests.Session for connection pooling
        """
        self.base_url = (
            base_url
            or os.getenv("AGENDEX_URL")
            or os.getenv("GOVERNANCE_URL")
            or "http://localhost:8000"
        ).rstrip("/")
        
        self.agent_id = (
            agent_id
            or os.getenv("AGENDEX_AGENT_ID")
            or os.getenv("GOVERNANCE_AGENT_ID")
        )
        self.tenant_id = (
            tenant_id
            or os.getenv("AGENDEX_TENANT_ID")
            or os.getenv("GOVERNANCE_TENANT_ID")
        )
        
        self.token = self._resolve_token(token)
        self.timeout = timeout
        self._session = session or requests.Session()
        self._mode = mode
        
        # Validate configuration
        if not self.agent_id:
            raise ConfigurationError(
                "agent_id is required. Set AGENDEX_AGENT_ID environment variable "
                "or pass agent_id to constructor."
            )
        if not self.token:
            raise ConfigurationError(
                "token is required. Set AGENDEX_TOKEN environment variable "
                "or pass token to constructor."
            )
        
        # Resolve mode from server if auto
        if self._mode == "auto":
            self._mode = self._fetch_mode()
        
        logger.info(
            "AgendexClient initialized: url=%s, agent=%s, mode=%s",
            self.base_url, self.agent_id, self._mode
        )
    
    @property
    def mode(self) -> str:
        """Current operating mode (shadow or enforce)."""
        return self._mode
    
    def _resolve_token(self, explicit: Optional[str]) -> Optional[str]:
        """Resolve token from explicit value or environment variables."""
        if explicit:
            return explicit
        
        # Try various env var patterns
        candidates = [
            os.getenv("AGENDEX_TOKEN"),
            os.getenv("AGENDEX_AGENT_TOKEN"),
            os.getenv("GOVERNANCE_AGENT_TOKEN"),
        ]
        
        # Try agent-specific token
        if self.agent_id:
            normalized = self.agent_id.upper().replace("-", "_")
            candidates.insert(0, os.getenv(f"{normalized}_TOKEN"))
            candidates.insert(1, os.getenv(f"TEST_{normalized}_TOKEN"))
        
        return next((v for v in candidates if v), None)
    
    def _headers(self) -> Dict[str, str]:
        """Build request headers with authentication."""
        headers: Dict[str, str] = {
            "Authorization": f"Bearer {self.token}",
            "Content-Type": "application/json",
            "X-Agendex-Agent-Id": self.agent_id,
            "X-Agendex-SDK-Version": "0.1.0",
        }
        if self.tenant_id:
            headers["X-Agendex-Tenant-Id"] = self.tenant_id
        return headers

    @staticmethod
    def _merge_context(
        context: Optional[Dict[str, Any]],
        *,
        trace_id: Optional[str] = None,
        span_id: Optional[str] = None,
        parent_span_id: Optional[str] = None,
        actor: Optional[str] = None,
    ) -> Optional[Dict[str, Any]]:
        merged = dict(context or {})
        if trace_id:
            merged["trace_id"] = trace_id
        if span_id:
            merged["span_id"] = span_id
        if parent_span_id:
            merged["parent_span_id"] = parent_span_id
        if actor:
            merged["actor"] = actor
        return merged or None
    
    def _fetch_mode(self) -> str:
        """Fetch current mode for this agent from server."""
        try:
            resp = self._session.get(
                f"{self.base_url}/agents/{self.agent_id}/mode",
                headers=self._headers(),
                timeout=5,
            )
            if resp.ok:
                return resp.json().get("mode", "shadow")
        except Exception as e:
            logger.debug(f"Could not fetch mode from server: {e}")
        
        # Default to shadow mode if can't reach server
        return "shadow"
    
    def invoke(
        self,
        action: str,
        params: Dict[str, Any],
        task: str,
        context: Optional[Dict[str, Any]] = None,
        resources: Optional[List[Dict[str, Any]]] = None,
        trace_id: Optional[str] = None,
        span_id: Optional[str] = None,
        parent_span_id: Optional[str] = None,
        actor: Optional[str] = None,
        on_event: Optional[Callable[[Dict], None]] = None,
    ) -> Dict[str, Any]:
        """
        Invoke an action through Agendex governance.
        
        In shadow mode: Action is allowed, activity is captured for analysis.
        In enforce mode: Action is evaluated against policy.
        
        Args:
            action: Action identifier (e.g., "tool.fetch_revenue", "http.request")
            params: Parameters for the action
            task: Task context grouping related actions
            context: Optional context (reasoning, user_prompt, step, etc.)
            resources: Optional explicit resources (auto-extracted if not provided)
            trace_id: Optional correlation trace identifier
            span_id: Optional current step span identifier
            parent_span_id: Optional parent step span identifier
            actor: Optional actor label (e.g., manager, assistant name)
            on_event: Optional callback for governance events
        
        Returns:
            In shadow mode: {"_agendex": "shadow", "decision": "allow"}
            In enforce mode: Result from the action execution
        
        Raises:
            DeniedError: If action is denied by policy
            PendingApprovalError: If action requires human approval
            AgendexError: For other errors
        """
        merged_context = self._merge_context(
            context,
            trace_id=trace_id,
            span_id=span_id,
            parent_span_id=parent_span_id,
            actor=actor,
        )

        # Build activity record
        activity = build_activity_record(
            agent_id=self.agent_id,
            action=action,
            params=params,
            task=task,
            context=merged_context,
            resources=resources,
        )
        
        # Emit tool_call event
        if on_event:
            on_event({
                "type": "tool_call",
                "action": action,
                "params": params,
                "task": task,
                "trace_id": trace_id,
                "span_id": span_id,
                "parent_span_id": parent_span_id,
                "actor": actor,
            })
        
        if self._mode == "shadow":
            return self._shadow_invoke(activity, on_event)
        else:
            return self._enforce_invoke(activity, on_event)
    
    def _shadow_invoke(
        self,
        activity: Dict[str, Any],
        on_event: Optional[Callable[[Dict], None]] = None,
    ) -> Dict[str, Any]:
        """Shadow mode: evaluate and emit what would happen, never invoke remotely."""
        # Log activity to server (best-effort)
        try:
            self._session.post(
                f"{self.base_url}/activity",
                json=activity,
                headers=self._headers(),
                timeout=5,
            )
        except Exception as e:
            logger.warning(f"Failed to log activity in shadow mode: {e}")

        decision = None
        request_id = None
        eval_latency_ms = None
        try:
            eval_start = perf_counter()
            eval_resp = self._session.post(
                f"{self.base_url}/evaluate",
                json={
                    "agent_id": activity["agent_id"],
                    "action": activity["action"],
                    "params": activity["params"],
                    "task": activity["task"],
                    "context": activity["context"],
                    "mode": "shadow",
                },
                headers=self._headers(),
                timeout=self.timeout,
            )
            if eval_resp.ok:
                eval_latency_ms = (perf_counter() - eval_start) * 1000.0
                eval_data = eval_resp.json()
                decision = eval_data.get("decision", {})
                request_id = eval_data.get("request_id")
                if isinstance(eval_data.get("latency_ms"), (int, float)):
                    eval_latency_ms = float(eval_data["latency_ms"])
        except Exception as e:
            logger.warning(f"Failed to evaluate in shadow mode: {e}")

        if on_event and decision:
            on_event({
                "type": "decision",
                "outcome": decision.get("outcome"),
                "reason": decision.get("reason"),
                "matched_rule": decision.get("matched_rule"),
                "mode": "shadow",
                "request_id": request_id,
                "latency_ms": eval_latency_ms,
                "action": activity["action"],
            })

        # Shadow returns governance metadata only; caller executes action locally.
        decision_outcome = decision.get("outcome") if isinstance(decision, dict) else None
        decision_reason = decision.get("reason") if isinstance(decision, dict) else None
        return {
            "_agendex": "shadow",
            "decision": decision_outcome or "unknown",
            "reason": decision_reason,
            "request_id": request_id,
            "execute": True,
        }
    
    def _enforce_invoke(
        self,
        activity: Dict[str, Any],
        on_event: Optional[Callable[[Dict], None]] = None,
    ) -> Dict[str, Any]:
        """Enforce mode: evaluate against policy and invoke if allowed."""
        # Step 1: Evaluate
        try:
            eval_resp = self._session.post(
                f"{self.base_url}/evaluate",
                json={
                    "agent_id": activity["agent_id"],
                    "action": activity["action"],
                    "params": activity["params"],
                    "task": activity["task"],
                    "context": activity["context"],
                },
                headers=self._headers(),
                timeout=self.timeout,
            )
            eval_resp.raise_for_status()
        except requests.exceptions.ConnectionError as e:
            raise ConnectionError(self.base_url, str(e)) from e
        except requests.exceptions.HTTPError as e:
            if eval_resp.status_code == 401:
                raise AgendexError(
                    f"Authentication failed: invalid token for agent '{self.agent_id}'"
                ) from e
            if eval_resp.status_code == 403:
                raise AgendexError(f"Authorization failed: {eval_resp.text}") from e
            raise AgendexError(f"Evaluate request failed: {e}") from e
        
        eval_data = eval_resp.json()
        decision = eval_data.get("decision", {})
        outcome = decision.get("outcome")
        reason = decision.get("reason", "No reason provided")
        matched_rule = decision.get("matched_rule")
        request_id = eval_data.get("request_id")
        latency_ms = eval_data.get("latency_ms")  # Latency is in response, not decision

        # Emit decision event
        if on_event:
            on_event({
                "type": "decision",
                "outcome": outcome,
                "reason": reason,
                "matched_rule": matched_rule,
                "mode": "enforce",
                "request_id": request_id,
                "latency_ms": latency_ms,
                "action": activity["action"],
            })
        
        # Step 2: Handle decision
        if outcome in ("deny", "would_deny"):
            raise DeniedError(
                action=activity["action"],
                reason=reason,
                matched_rule=matched_rule,
                request_id=request_id,
            )
        
        if outcome == "pending_approval":
            approval_id = eval_data.get("approval_id", "unknown")
            if on_event:
                on_event({
                    "type": "approval_requested",
                    "approval_id": approval_id,
                    "action": activity["action"],
                    "params": activity["params"],
                    "reason": reason,
                    "matched_rule": matched_rule,
                })
            raise PendingApprovalError(
                action=activity["action"],
                approval_id=approval_id,
                reason=reason,
                matched_rule=matched_rule,
            )
        
        if outcome == "would_allow":
            raise ShadowDecisionError(
                action=activity["action"],
                outcome=outcome,
                reason=reason,
            )
        
        if outcome != "allow":
            raise AgendexError(f"Unexpected decision outcome: {outcome}")
        
        # Step 3: Invoke with scoped token
        scoped_token = eval_data.get("token")
        if not scoped_token:
            raise AgendexError("Evaluation allowed but no token returned")
        
        try:
            invoke_start = perf_counter()
            invoke_resp = self._session.post(
                f"{self.base_url}/invoke",
                json={
                    "token": scoped_token,
                    "action": activity["action"],
                    "params": activity["params"],
                },
                headers=self._headers(),
                timeout=self.timeout,
            )
            invoke_resp.raise_for_status()
            invoke_latency_ms = (perf_counter() - invoke_start) * 1000.0
        except requests.exceptions.HTTPError as e:
            if invoke_resp.status_code == 403:
                raise AgendexError(f"Invoke forbidden: {invoke_resp.text}") from e
            raise AgendexError(f"Invoke request failed: {e}") from e
        
        invoke_data = invoke_resp.json()
        
        if not invoke_data.get("success"):
            raise AgendexError(f"Invoke failed: {invoke_data}")
        
        result = invoke_data.get("result", {})
        
        # Emit result event
        if on_event:
            on_event({
                "type": "invoke_result",
                "result": result,
                "latency_ms": invoke_latency_ms,
            })
        
        return result
    
    def check_approval(self, approval_id: str) -> Dict[str, Any]:
        """
        Check the status of a pending approval.
        
        Args:
            approval_id: The approval ID from PendingApprovalError
        
        Returns:
            Approval status dict with 'status', 'token' (if approved), etc.
        """
        try:
            resp = self._session.get(
                f"{self.base_url}/approvals",
                headers=self._headers(),
                timeout=self.timeout,
            )
            resp.raise_for_status()
        except requests.exceptions.RequestException as e:
            raise AgendexError(f"Failed to check approvals: {e}") from e
        
        approvals = resp.json().get("approvals", [])
        for approval in approvals:
            if approval.get("approval_id") == approval_id:
                return approval
        
        raise AgendexError(f"Approval '{approval_id}' not found")

    def ingest_events(self, events: List[Dict[str, Any]]) -> Dict[str, Any]:
        """
        Persist a batch of trace events to the governance proxy.

        Args:
            events: List of event dicts

        Returns:
            Server response, typically {"ingested": N}
        """
        payload_events = [event for event in events if isinstance(event, dict)]
        if not payload_events:
            return {"ingested": 0}

        try:
            resp = self._session.post(
                f"{self.base_url}/events/ingest",
                json={"events": payload_events},
                headers=self._headers(),
                timeout=min(self.timeout, 10),
            )
            resp.raise_for_status()
            data = resp.json()
            return data if isinstance(data, dict) else {"ingested": len(payload_events)}
        except requests.exceptions.RequestException as e:
            raise AgendexError(f"Failed to ingest events: {e}") from e
    
    def set_mode(self, mode: str) -> None:
        """
        Request mode change for this agent.
        
        Args:
            mode: New mode ("shadow" or "enforce")
        """
        if mode not in ("shadow", "enforce"):
            raise ValueError(f"Invalid mode: {mode}. Must be 'shadow' or 'enforce'")
        
        try:
            resp = self._session.put(
                f"{self.base_url}/agents/{self.agent_id}/mode",
                json={"mode": mode},
                headers=self._headers(),
                timeout=10,
            )
            resp.raise_for_status()
            self._mode = mode
            logger.info(f"Mode changed to: {mode}")
        except Exception as e:
            raise AgendexError(f"Failed to set mode: {e}") from e
