"""
StrivioClient — lightweight HTTP client for the Strivio AI Firewall API.
"""
from __future__ import annotations

import os
import time
import uuid
import logging
import threading
from typing import Any, Dict, Optional

try:
    import requests
    _HAS_REQUESTS = True
except ImportError:
    _HAS_REQUESTS = False

try:
    import urllib.request
    import urllib.error
    import json as _json
    _HAS_URLLIB = True
except ImportError:
    _HAS_URLLIB = False

logger = logging.getLogger("strivio")


class StrivioClient:
    """
    Sends agent events to the Strivio AI Firewall.

    Usage:
        client = StrivioClient(
            api_url="http://localhost:4000/api/v1",
            token="your-jwt-token",
        )

        result = client.send_event(
            agent_id="my-langchain-agent",
            action="tool_call",
            target="search_api",
            risk_score=45,
            metadata={"tool": "SerpAPI", "query": "quarterly earnings"},
        )
        print(result)  # {"firewallAction": "allow", "riskScore": 45, ...}
    """

    def __init__(
        self,
        api_url:    Optional[str] = None,
        token:      Optional[str] = None,
        agent_id:   Optional[str] = None,
        async_mode: bool          = False,
        timeout:    int           = 15,
    ):
        self.api_url    = (api_url or os.getenv("STRIVIO_API_URL", "http://localhost:4000/api/v1")).rstrip("/")
        self.token      = token  or os.getenv("STRIVIO_TOKEN", "")
        self.agent_id   = agent_id or os.getenv("STRIVIO_AGENT_ID", "unknown-agent")
        self.async_mode = async_mode
        self.timeout    = timeout

        if not self.token:
            logger.warning("[Strivio] No token set — events will be rejected by the firewall.")

    # ── Public API ────────────────────────────────────────────────────────

    def send_event(
        self,
        action:     str,
        target:     str,
        agent_id:   Optional[str]       = None,
        risk_score: Optional[int]       = None,
        metadata:   Optional[Dict]      = None,
        event_type: str                 = "tool_call",
    ) -> Optional[Dict]:
        """
        Send a single agent event to Strivio.

        Args:
            action:     What the agent did    (e.g. "tool_call", "data_export", "api_call")
            target:     What it acted on      (e.g. "search_api", "s3://bucket", "send_email")
            agent_id:   Override the client-level agent_id for this event
            risk_score: Manual risk hint 0–100 (Strivio will re-score internally)
            metadata:   Arbitrary dict attached to the audit log
            event_type: High-level category (default "tool_call")

        Returns:
            Dict with firewallAction, riskScore, blocked, eventId — or None on error.
        """
        payload = {
            "agentId":   agent_id or self.agent_id,
            "action":    action,
            "target":    target,
            "eventType": event_type,
            "metadata":  metadata or {},
        }
        if risk_score is not None:
            payload["riskScore"] = risk_score

        if self.async_mode:
            t = threading.Thread(target=self._post, args=(payload,), daemon=True)
            t.start()
            return None
        return self._post(payload)

    def tool_call(self, tool_name: str, tool_input: Any, agent_id: Optional[str] = None) -> Optional[Dict]:
        """Shortcut for tool invocation events."""
        return self.send_event(
            action="tool_call",
            target=tool_name,
            agent_id=agent_id,
            metadata={"tool": tool_name, "input": str(tool_input)[:500]},
        )

    def llm_call(self, prompt: str, model: str = "unknown", agent_id: Optional[str] = None) -> Optional[Dict]:
        """Shortcut for LLM call events."""
        return self.send_event(
            action="llm_call",
            target=model,
            agent_id=agent_id,
            metadata={"model": model, "prompt_preview": prompt[:300]},
        )

    def data_access(self, source: str, record_count: int = 0, agent_id: Optional[str] = None) -> Optional[Dict]:
        """Shortcut for data access events."""
        return self.send_event(
            action="data_access",
            target=source,
            agent_id=agent_id,
            risk_score=60 if record_count > 100 else 30,
            metadata={"source": source, "record_count": record_count},
        )

    def is_blocked(self, result: Optional[Dict]) -> bool:
        """Returns True if Strivio blocked the action."""
        if result is None:
            return False
        return result.get("blocked", False) or result.get("firewallAction") == "block"

    # ── Internal ──────────────────────────────────────────────────────────

    def _post(self, payload: Dict) -> Optional[Dict]:
        url     = f"{self.api_url}/events"
        headers = {
            "Content-Type":  "application/json",
            "Authorization": f"Bearer {self.token}",
        }

        try:
            if _HAS_REQUESTS:
                resp = requests.post(url, json=payload, headers=headers, timeout=self.timeout)
                data = resp.json()
            elif _HAS_URLLIB:
                import json as _json
                req  = urllib.request.Request(
                    url,
                    data=_json.dumps(payload).encode(),
                    headers=headers,
                    method="POST",
                )
                with urllib.request.urlopen(req, timeout=self.timeout) as r:
                    data = _json.loads(r.read())
            else:
                logger.error("[Strivio] No HTTP library available (install requests)")
                return None

            if data.get("success"):
                result = data.get("data", {})
                if result.get("blocked"):
                    logger.warning(
                        "[Strivio] BLOCKED  agent=%s action=%s target=%s score=%s",
                        payload.get("agentId"), payload.get("action"),
                        payload.get("target"), result.get("riskScore"),
                    )
                else:
                    logger.debug(
                        "[Strivio] ALLOWED  agent=%s action=%s score=%s",
                        payload.get("agentId"), payload.get("action"), result.get("riskScore"),
                    )
                return result
            else:
                logger.error("[Strivio] API error: %s", data.get("message"))
                return None

        except Exception as e:
            logger.error("[Strivio] Request failed: %s", e)
            return None
