#!/usr/bin/env python3
"""
Bloom Secure Proxy Skill for OpenClaw

Routes all outbound API requests through Bloom's IAM proxy for
authentication, authorization, audit logging, and security scanning.
"""

import os
import json
import hashlib
import time
from typing import Optional, Dict, Any, List
from urllib.request import Request, urlopen
from urllib.error import HTTPError, URLError

# Configuration
BLOOM_PROXY_URL = os.environ.get(
    "BLOOM_PROXY_URL",
    "https://iam.bloomtechnologies.app"
)
BLOOM_AGENT_TOKEN = os.environ.get("BLOOM_AGENT_TOKEN")
BLOOM_ORG_ID = os.environ.get("BLOOM_ORG_ID")

# Timeouts
REQUEST_TIMEOUT = int(os.environ.get("BLOOM_TIMEOUT", "30"))
MAX_RETRIES = int(os.environ.get("BLOOM_RETRIES", "3"))

# Cache (optional, in-memory)
_response_cache: Dict[str, tuple] = {}
CACHE_ENABLED = os.environ.get("BLOOM_CACHE_ENABLED", "false").lower() == "true"
CACHE_TTL = int(os.environ.get("BLOOM_CACHE_TTL", "300"))


class BloomProxyError(Exception):
    """Raised when Bloom proxy returns an error."""
    def __init__(self, status_code: int, message: str, details: Optional[dict] = None):
        self.status_code = status_code
        self.message = message
        self.details = details or {}
        super().__init__(f"Bloom Proxy Error {status_code}: {message}")


class BloomSecurityBlock(Exception):
    """Raised when Bloom blocks a request for security reasons."""
    def __init__(self, reason: str, threat_type: str):
        self.reason = reason
        self.threat_type = threat_type
        super().__init__(f"Request blocked: {reason} (threat: {threat_type})")


class BloomKillSwitchActive(Exception):
    """Raised when the agent has been killed."""
    def __init__(self, agent_id: str, killed_at: str):
        self.agent_id = agent_id
        self.killed_at = killed_at
        super().__init__(f"Agent {agent_id} has been killed at {killed_at}")


def _get_cache_key(method: str, url: str, body: Optional[str]) -> str:
    """Generate cache key for a request."""
    content = f"{method}:{url}:{body or ''}"
    return hashlib.sha256(content.encode()).hexdigest()


def _check_cache(cache_key: str) -> Optional[dict]:
    """Check if response is cached and not expired."""
    if not CACHE_ENABLED:
        return None

    if cache_key in _response_cache:
        response, timestamp = _response_cache[cache_key]
        if time.time() - timestamp < CACHE_TTL:
            return response
        else:
            del _response_cache[cache_key]
    return None


def _store_cache(cache_key: str, response: dict):
    """Store response in cache."""
    if CACHE_ENABLED:
        _response_cache[cache_key] = (response, time.time())


def bloom_request(
    method: str,
    url: str,
    headers: Optional[Dict[str, str]] = None,
    body: Optional[Any] = None,
    timeout: Optional[int] = None,
    skip_cache: bool = False
) -> Dict[str, Any]:
    """
    Route an HTTP request through Bloom's secure proxy.

    Args:
        method: HTTP method (GET, POST, PUT, DELETE, etc.)
        url: Target URL
        headers: Optional headers to include
        body: Optional request body (will be JSON-encoded if dict)
        timeout: Optional timeout override
        skip_cache: Skip cache lookup/storage

    Returns:
        dict with 'status_code', 'headers', 'body' keys

    Raises:
        BloomProxyError: If proxy returns an error
        BloomSecurityBlock: If request is blocked for security
        BloomKillSwitchActive: If agent has been killed
    """
    if not BLOOM_AGENT_TOKEN:
        raise BloomProxyError(
            401,
            "BLOOM_AGENT_TOKEN not set. Please configure your Bloom credentials."
        )

    # Serialize body
    body_str = None
    if body is not None:
        if isinstance(body, (dict, list)):
            body_str = json.dumps(body)
        else:
            body_str = str(body)

    # Check cache
    cache_key = _get_cache_key(method, url, body_str)
    if not skip_cache:
        cached = _check_cache(cache_key)
        if cached:
            return cached

    # Build proxy URL - use the direct pass-through format
    # e.g., https://iam.bloomtechnologies.app/https://api.openai.com/v1/chat/completions
    proxy_url = f"{BLOOM_PROXY_URL}/{url}"

    proxy_headers = {
        "Content-Type": headers.get("Content-Type", "application/json") if headers else "application/json",
        "Authorization": f"Bearer {BLOOM_AGENT_TOKEN}",
        "X-Bloom-Agent-Version": "1.0.0",
        "X-Bloom-Skill": "bloom-openclaw-skill"
    }

    # Forward original headers (except Authorization which Bloom will inject)
    if headers:
        for key, value in headers.items():
            if key.lower() not in ["authorization", "content-type"]:
                proxy_headers[key] = value

    if BLOOM_ORG_ID:
        proxy_headers["X-Bloom-Org-ID"] = BLOOM_ORG_ID

    # Make request with retries
    last_error = None
    for attempt in range(MAX_RETRIES):
        try:
            req = Request(
                proxy_url,
                data=body_str.encode("utf-8") if body_str else None,
                headers=proxy_headers,
                method=method.upper()
            )

            with urlopen(req, timeout=timeout or REQUEST_TIMEOUT) as resp:
                response_body = resp.read().decode("utf-8")

                result = {
                    "status_code": resp.status,
                    "headers": dict(resp.headers),
                    "body": response_body
                }

                # Try to parse JSON response
                try:
                    result["json"] = json.loads(response_body)
                except json.JSONDecodeError:
                    result["json"] = None

                # Cache successful GET responses
                if method.upper() == "GET" and result["status_code"] == 200:
                    _store_cache(cache_key, result)

                return result

        except HTTPError as e:
            error_body = e.read().decode("utf-8")
            try:
                error_data = json.loads(error_body)
            except json.JSONDecodeError:
                error_data = {"message": error_body}

            # Check for kill switch
            if e.code == 403 and error_data.get("error") == "Agent killed":
                raise BloomKillSwitchActive(
                    error_data.get("agent_id", "unknown"),
                    error_data.get("killed_at", "unknown")
                )

            # Check for security block
            if e.code == 403 and error_data.get("security_block"):
                raise BloomSecurityBlock(
                    error_data.get("reason", "Request blocked"),
                    error_data.get("threat_type", "unknown")
                )

            # Don't retry 4xx errors (except 429)
            if 400 <= e.code < 500 and e.code != 429:
                raise BloomProxyError(e.code, error_data.get("message", str(e)), error_data)

            last_error = BloomProxyError(e.code, error_data.get("message", str(e)), error_data)

        except URLError as e:
            last_error = BloomProxyError(0, f"Network error: {e.reason}")

        # Exponential backoff
        if attempt < MAX_RETRIES - 1:
            time.sleep(2 ** attempt)

    raise last_error


def bloom_get(url: str, headers: Optional[Dict[str, str]] = None, **kwargs) -> Dict[str, Any]:
    """Convenience wrapper for GET requests."""
    return bloom_request("GET", url, headers=headers, **kwargs)


def bloom_post(url: str, body: Any, headers: Optional[Dict[str, str]] = None, **kwargs) -> Dict[str, Any]:
    """Convenience wrapper for POST requests."""
    return bloom_request("POST", url, headers=headers, body=body, **kwargs)


def bloom_put(url: str, body: Any, headers: Optional[Dict[str, str]] = None, **kwargs) -> Dict[str, Any]:
    """Convenience wrapper for PUT requests."""
    return bloom_request("PUT", url, headers=headers, body=body, **kwargs)


def bloom_delete(url: str, headers: Optional[Dict[str, str]] = None, **kwargs) -> Dict[str, Any]:
    """Convenience wrapper for DELETE requests."""
    return bloom_request("DELETE", url, headers=headers, **kwargs)


def bloom_patch(url: str, body: Any, headers: Optional[Dict[str, str]] = None, **kwargs) -> Dict[str, Any]:
    """Convenience wrapper for PATCH requests."""
    return bloom_request("PATCH", url, headers=headers, body=body, **kwargs)


# Health check
def check_bloom_connection() -> Dict[str, Any]:
    """Verify connection to Bloom proxy and get agent status."""
    try:
        req = Request(
            f"{BLOOM_PROXY_URL}/health",
            headers={"Authorization": f"Bearer {BLOOM_AGENT_TOKEN}"} if BLOOM_AGENT_TOKEN else {}
        )
        with urlopen(req, timeout=5) as resp:
            data = json.loads(resp.read().decode("utf-8"))
            return {
                "connected": True,
                "proxy_status": data.get("status", "unknown"),
                "proxy_url": BLOOM_PROXY_URL,
                "agent_token_set": bool(BLOOM_AGENT_TOKEN)
            }
    except Exception as e:
        return {
            "connected": False,
            "error": str(e),
            "proxy_url": BLOOM_PROXY_URL,
            "agent_token_set": bool(BLOOM_AGENT_TOKEN)
        }


def get_agent_status() -> Dict[str, Any]:
    """Get the current agent's status from Bloom."""
    if not BLOOM_AGENT_TOKEN:
        return {"error": "BLOOM_AGENT_TOKEN not set"}

    try:
        # Extract agent ID from token if it's in the format bloom_xxx_agent_yyy
        agent_id = None
        if "_agent_" in BLOOM_AGENT_TOKEN:
            agent_id = BLOOM_AGENT_TOKEN.split("_agent_")[1]

        req = Request(
            f"{BLOOM_PROXY_URL}/health",
            headers={"Authorization": f"Bearer {BLOOM_AGENT_TOKEN}"}
        )
        with urlopen(req, timeout=5) as resp:
            return {
                "status": "active",
                "agent_id": agent_id,
                "proxy_reachable": True
            }
    except HTTPError as e:
        if e.code == 403:
            return {
                "status": "killed",
                "agent_id": agent_id if 'agent_id' in dir() else None,
                "proxy_reachable": True
            }
        return {
            "status": "error",
            "error": str(e),
            "proxy_reachable": True
        }
    except Exception as e:
        return {
            "status": "error",
            "error": str(e),
            "proxy_reachable": False
        }


if __name__ == "__main__":
    # Self-test
    print("Bloom Secure Proxy - Connection Test")
    print("=" * 40)

    status = check_bloom_connection()

    if status["connected"]:
        print(f"  Connected: Yes")
        print(f"  Proxy URL: {status['proxy_url']}")
        print(f"  Proxy Status: {status['proxy_status']}")
        print(f"  Agent Token Set: {status['agent_token_set']}")
    else:
        print(f"  Connected: No")
        print(f"  Error: {status.get('error', 'Unknown')}")
        print(f"  Proxy URL: {status['proxy_url']}")
        print(f"  Agent Token Set: {status['agent_token_set']}")

    print()

    if not status['agent_token_set']:
        print("Set BLOOM_AGENT_TOKEN environment variable to enable proxy.")
        print("Get your token at: https://platform.bloomtechnologies.app/dashboard")
