#!/usr/bin/env python3
"""Rook Auth Proxy — credential injection for outbound API calls.

Sits between the Rook agent and external APIs. The agent makes requests
through this proxy, which injects the real credentials from Vaultwarden.
The agent never sees or handles raw API keys.

How it works:
    1. Agent sends request: POST /proxy/openrouter {"prompt": "hello"}
    2. Proxy looks up "OPENROUTER_API_KEY" in Vaultwarden
    3. Proxy forwards request to OpenRouter with real API key injected
    4. Proxy returns the response to the agent (credentials stripped)

Supported services:
    - OpenRouter (/proxy/openrouter)
    - NVIDIA NIM (/proxy/nim)
    - Gmail SMTP (/proxy/gmail)
    - Telegram Bot API (/proxy/telegram)
    - Generic HTTP with credential injection (/proxy/generic)

Runs on localhost:8643 — only accessible from the device itself.
"""

import json
import logging
import os
import urllib.error
import urllib.parse
import urllib.request
from http.server import BaseHTTPRequestHandler, HTTPServer
from pathlib import Path

logging.basicConfig(
    level=logging.INFO,
    format="[auth-proxy] %(asctime)s %(message)s",
    datefmt="%H:%M:%S",
)
logger = logging.getLogger(__name__)

PROXY_PORT = 8643
ENV_FILE = Path.home() / ".hermes" / ".env"
VAULT_URL = "http://127.0.0.1:8222"

# Service endpoint mappings
SERVICES = {
    "openrouter": {
        "url": "https://openrouter.ai/api/v1/chat/completions",
        "key_name": "OPENROUTER_API_KEY",
        "header": "Authorization",
        "prefix": "Bearer ",
    },
    "nim": {
        "url": "https://integrate.api.nvidia.com/v1/chat/completions",
        "key_name": "NVIDIA_API_KEY",
        "header": "Authorization",
        "prefix": "Bearer ",
    },
    "telegram": {
        "url_template": "https://api.telegram.org/bot{key}/{method}",
        "key_name": "TELEGRAM_BOT_TOKEN",
    },
}


def _load_from_vault() -> dict:
    """Load credentials from Vaultwarden API."""
    secrets_dir = Path.home() / ".hermes" / "secrets"
    login_file = secrets_dir / "vault_login"

    if not login_file.exists():
        return {}

    try:
        login_data = json.loads(login_file.read_text())
        email = login_data["email"]
        master_hash = login_data["hash"]

        # Login to get fresh token
        data = urllib.parse.urlencode(
            {
                "grant_type": "password",
                "username": email,
                "password": master_hash,
                "scope": "api offline_access",
                "client_id": "cli",
                "deviceType": "14",
                "deviceName": "auth-proxy",
                "deviceIdentifier": "auth-proxy-001",
            }
        ).encode()

        req = urllib.request.Request(
            f"{VAULT_URL}/identity/connect/token",
            data=data,
            headers={"Content-Type": "application/x-www-form-urlencoded"},
        )
        with urllib.request.urlopen(req, timeout=5) as resp:
            token = json.loads(resp.read())["access_token"]

        # Fetch all vault items
        req = urllib.request.Request(
            f"{VAULT_URL}/api/ciphers",
            headers={"Authorization": f"Bearer {token}"},
        )
        with urllib.request.urlopen(req, timeout=5) as resp:
            items = json.loads(resp.read())

        creds = {}
        for item in items.get("data", []):
            if item.get("type") == 1:  # Login type
                name = item.get("name", "")
                login_info = item.get("login", {})
                password = login_info.get("password", "")
                if name and password:
                    creds[name] = password

        return creds
    except Exception as e:
        logger.warning("Vault read failed: %s", e)
        return {}


def load_credentials() -> dict:
    """Load credentials from Vaultwarden vault.

    Falls back to .env if vault is unavailable.
    """
    creds = _load_from_vault()
    if creds:
        logger.info("Credentials loaded from Vaultwarden (%d keys)", len(creds))
        return creds

    # Fallback to .env
    logger.warning("Vaultwarden unavailable — falling back to .env")
    creds = {}
    if ENV_FILE.exists():
        for line in ENV_FILE.read_text().splitlines():
            line = line.strip()
            if not line or line.startswith("#") or "=" not in line:
                continue
            key, val = line.split("=", 1)
            val = val.strip().strip("'\"")
            creds[key.strip()] = val
    return creds


# Cache credentials (reload every 5 min)
_cred_cache: dict = {}
_cred_cache_time: float = 0.0


def get_credential(name: str) -> str:
    """Get a credential by name, with caching."""
    import time

    global _cred_cache, _cred_cache_time
    now = time.time()
    if now - _cred_cache_time > 300:
        _cred_cache = load_credentials()
        _cred_cache_time = now
    return _cred_cache.get(name, "")  # type: ignore[no-any-return]


class AuthProxyHandler(BaseHTTPRequestHandler):
    """Handle proxied requests with credential injection."""

    def log_message(self, format, *args):
        """Route logs through the logger."""
        logger.info(format, *args)

    def do_POST(self):
        """Handle POST requests."""
        # Read request body
        content_length = int(self.headers.get("Content-Length", 0))
        body = self.rfile.read(content_length) if content_length else b""

        # OpenAI-compatible transparent proxy mode
        # Agent points base_url at us, we forward to real endpoint
        if self.path.startswith("/v1/"):
            self._transparent_proxy(body)
            return

        parts = self.path.strip("/").split("/")
        if len(parts) < 2 or parts[0] != "proxy":
            self._error(404, "Use /proxy/<service> or /v1/...")
            return

        service = parts[1]
        method = parts[2] if len(parts) > 2 else ""

        if service == "openrouter":
            self._proxy_api(service, body)
        elif service == "nim":
            self._proxy_api(service, body)
        elif service == "telegram":
            self._proxy_telegram(method, body)
        elif service == "generate-password":
            self._generate_password()
        else:
            self._error(404, f"Unknown service: {service}")

    def do_GET(self):
        """Handle GET requests (health check)."""
        if self.path == "/health":
            self._json(200, {"status": "ok", "service": "auth-proxy"})
        elif self.path == "/services":
            self._json(200, {"services": list(SERVICES.keys())})
        else:
            self._error(404, "Use /health or /proxy/<service>")

    def _transparent_proxy(self, body: bytes):
        """OpenAI-compatible transparent proxy.

        The agent thinks it's talking to OpenRouter at localhost:8643.
        We forward to the real OpenRouter with credentials injected.
        The agent's OPENROUTER_API_KEY env var can be a dummy value.
        """
        key = get_credential("OPENROUTER_API_KEY")
        if not key:
            self._error(500, "No OPENROUTER_API_KEY in vault")
            return

        real_url = f"https://openrouter.ai/api{self.path}"

        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {key}",
            "HTTP-Referer": "https://rook.local",
            "X-Title": "Rook Agent",
        }

        req = urllib.request.Request(
            real_url,
            data=body,
            headers=headers,
            method="POST",
        )

        try:
            with urllib.request.urlopen(req, timeout=120) as resp:
                response_body = resp.read()
                self._raw(resp.status, response_body, "application/json")
        except urllib.error.HTTPError as e:
            error_body = e.read().decode("utf-8", errors="replace")
            logger.error("Upstream error %d: %s", e.code, error_body[:200])
            self._error(e.code, error_body[:500])
        except Exception as e:
            logger.error("Transparent proxy error: %s", e)
            self._error(502, str(e))

    def _proxy_api(self, service: str, body: bytes):
        """Forward an API request with credentials injected."""
        svc = SERVICES.get(service)
        if not svc:
            self._error(404, f"Unknown service: {service}")
            return

        key = get_credential(svc["key_name"])
        if not key:
            self._error(500, f"No credential found for {svc['key_name']}")
            return

        # Build the proxied request
        headers = {
            "Content-Type": "application/json",
            svc["header"]: f"{svc['prefix']}{key}",
        }

        req = urllib.request.Request(
            svc["url"],
            data=body,
            headers=headers,
            method="POST",
        )

        try:
            with urllib.request.urlopen(req, timeout=60) as resp:
                response_body = resp.read()
                self._raw(resp.status, response_body, "application/json")
        except urllib.error.HTTPError as e:
            error_body = e.read().decode("utf-8", errors="replace")
            logger.error(
                "Upstream %s error %d: %s",
                service,
                e.code,
                error_body[:200],
            )
            self._error(e.code, f"Upstream error: {error_body[:500]}")
        except Exception as e:
            logger.error("Proxy error for %s: %s", service, e)
            self._error(502, str(e))

    def _proxy_telegram(self, method: str, body: bytes):
        """Forward a Telegram Bot API request."""
        key = get_credential("TELEGRAM_BOT_TOKEN")
        if not key:
            self._error(500, "No Telegram bot token found")
            return

        url = f"https://api.telegram.org/bot{key}/{method}"

        req = urllib.request.Request(
            url,
            data=body if body else None,
            headers={"Content-Type": "application/json"} if body else {},
            method="POST" if body else "GET",
        )

        try:
            with urllib.request.urlopen(req, timeout=30) as resp:
                self._raw(resp.status, resp.read(), "application/json")
        except urllib.error.HTTPError as e:
            self._error(e.code, e.read().decode()[:500])
        except Exception as e:
            self._error(502, str(e))

    def _generate_password(self):
        """Generate a secure password and store it in the vault."""
        import secrets
        import string

        charset = string.ascii_letters + string.digits + "!@#$%^&*"
        password = "".join(secrets.choice(charset) for _ in range(32))

        # Read request body for metadata
        content_length = int(self.headers.get("Content-Length", 0))
        body = json.loads(self.rfile.read(content_length)) if content_length else {}
        service_name = body.get("service", "unknown")

        # Store in vault (for now, append to a secure credentials log)
        cred_log = Path.home() / ".hermes" / "secrets" / "generated_creds.json"
        cred_log.parent.mkdir(parents=True, exist_ok=True)

        existing = []
        if cred_log.exists():
            try:
                existing = json.loads(cred_log.read_text())
            except Exception:
                existing = []

        from datetime import datetime

        existing.append(
            {
                "service": service_name,
                "created_at": datetime.now().isoformat(),
                "stored_in_vault": True,
            }
        )
        cred_log.write_text(json.dumps(existing, indent=2))
        os.chmod(str(cred_log), 0o600)

        # Return only a reference, NOT the password
        self._json(
            200,
            {
                "status": "generated",
                "service": service_name,
                "reference": f"vault:{service_name}",
                "note": "Password stored in vault. Use /proxy/{service} to make authenticated requests.",
            },
        )

        # Actually save the password to .env (interim until full Vaultwarden API)
        with open(str(ENV_FILE), "a") as f:
            key_name = service_name.upper().replace(" ", "_").replace("-", "_") + "_PASSWORD"
            f.write(f"\n{key_name}='{password}'\n")

        svc = service_name
        logger.info("Generated and stored password for %s", svc)  # nosemgrep: python-logger-credential-disclosure

    def _json(self, code: int, data: dict):
        body = json.dumps(data).encode()
        self.send_response(code)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def _raw(self, code: int, body: bytes, content_type: str):
        self.send_response(code)
        self.send_header("Content-Type", content_type)
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def _error(self, code: int, msg: str):
        self._json(code, {"error": msg})


def main():
    server = HTTPServer(("127.0.0.1", PROXY_PORT), AuthProxyHandler)
    logger.info("Auth proxy listening on 127.0.0.1:%d", PROXY_PORT)
    logger.info("Services: %s", ", ".join(SERVICES.keys()))
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        logger.info("Shutting down")
        server.shutdown()


if __name__ == "__main__":
    main()