"""
OpenAI-compatible proxy adapter for llm-guard-kit.

Intercepts OpenAI /v1/chat/completions (and /v1/completions) requests,
forwards them to an upstream LLM, then reconstructs a pseudo-ReAct chain
from the response's tool_calls / function_call fields and pipes it through
the llm-guard-kit HTTP API for reliability scoring.

Compatible with:
  - Open WebUI   (set OpenAI base URL to http://localhost:8001)
  - LibreChat    (same)
  - Chatbot UI   (same)
  - Any OpenAI-SDK client

Usage
-----
    # Start llm-guard HTTP API on :8000 first
    llm-guard-kit serve --domain my_domain --port 8000

    # Start this proxy on :8001
    python -m qppg_service.openai_proxy --upstream https://api.openai.com \
           --guard-url http://localhost:8000 --port 8001

    # In Open WebUI → Settings → OpenAI → set base URL to http://localhost:8001

Environment variables (override CLI flags):
    UPSTREAM_BASE_URL   upstream OpenAI-compatible base (default: https://api.openai.com)
    UPSTREAM_API_KEY    API key forwarded to upstream
    GUARD_URL           llm-guard-kit HTTP API URL (default: http://localhost:8000)
    GUARD_THRESHOLD     risk threshold for X-LLMGuard-Alert header (default: 0.65)
    PROXY_PORT          port to listen on (default: 8001)

Design
------
The proxy adds two headers to every response:
    X-LLMGuard-Risk:  float  0.0–1.0 chain risk score
    X-LLMGuard-Alert: bool   true if risk > threshold

For tool_call-based responses the proxy reconstructs a ReAct chain:
    each tool call → Step(action_type=fn_name, action_arg=args_json,
                          observation=<tool_result if available else "">)
For plain text responses (no tool calls) the proxy creates a minimal
single-step chain with the full text as the final answer.

The guard score is added to the response body under
    choices[0].message.guard  (non-breaking extension field)
so clients that ignore unknown fields are unaffected.
"""

from __future__ import annotations

import argparse
import json
import os
import time
from typing import Any, Dict, List, Optional

import httpx
from fastapi import FastAPI, HTTPException, Request, Response
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, StreamingResponse
from pydantic import BaseModel

# ── Config ────────────────────────────────────────────────────────────────────

DEFAULT_UPSTREAM  = "https://api.openai.com"
DEFAULT_GUARD_URL = "http://localhost:8000"
DEFAULT_THRESHOLD = 0.65
DEFAULT_PORT      = 8001

# ── App ───────────────────────────────────────────────────────────────────────

app = FastAPI(
    title="llm-guard-kit OpenAI proxy",
    description="OpenAI-compatible proxy that scores responses via llm-guard-kit",
    version="1.0.0",
)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
    expose_headers=["X-LLMGuard-Risk", "X-LLMGuard-Alert"],
)

# Runtime config — populated in main() or via env vars
_config: Dict[str, Any] = {}


# ── Chain reconstruction ──────────────────────────────────────────────────────

def _tool_calls_to_steps(tool_calls: List[Dict]) -> List[Dict]:
    """Convert OpenAI tool_calls array to llm-guard-kit Step list."""
    steps = []
    for tc in tool_calls:
        fn   = tc.get("function", {})
        name = fn.get("name", "tool_call")
        args = fn.get("arguments", "{}")
        # Try to pretty-print the args as a flat string for action_arg
        try:
            parsed = json.loads(args)
            action_arg = " ".join(str(v) for v in parsed.values()) if parsed else args
        except (json.JSONDecodeError, AttributeError):
            action_arg = str(args)
        steps.append({
            "thought":     f"Calling {name}",
            "action_type": name,
            "action_arg":  action_arg,
            "observation": "",  # upstream doesn't return tool results inline
        })
    return steps


def _function_call_to_steps(function_call: Dict) -> List[Dict]:
    """Convert legacy function_call dict to a single Step."""
    name = function_call.get("name", "function_call")
    args = function_call.get("arguments", "{}")
    try:
        parsed = json.loads(args)
        action_arg = " ".join(str(v) for v in parsed.values()) if parsed else args
    except (json.JSONDecodeError, AttributeError):
        action_arg = str(args)
    return [{"thought": f"Calling {name}", "action_type": name,
             "action_arg": action_arg, "observation": ""}]


def _build_chain(question: str, message: Dict, final_answer: str) -> Dict:
    """Reconstruct a pseudo-ReAct chain from a chat completion message."""
    tool_calls    = message.get("tool_calls") or []
    function_call = message.get("function_call")

    if tool_calls:
        steps = _tool_calls_to_steps(tool_calls)
        # Add a Finish step with the text content (if any) as the final answer
        steps.append({
            "thought":     "Returning final answer.",
            "action_type": "Finish",
            "action_arg":  final_answer,
            "observation": "",
        })
    elif function_call:
        steps = _function_call_to_steps(function_call)
        steps.append({
            "thought":     "Returning final answer.",
            "action_type": "Finish",
            "action_arg":  final_answer,
            "observation": "",
        })
    else:
        # Plain text response — single-step chain
        steps = [{
            "thought":     "Generating answer.",
            "action_type": "Finish",
            "action_arg":  final_answer,
            "observation": "",
        }]

    return {"question": question, "steps": steps, "final_answer": final_answer}


# ── Guard scoring ─────────────────────────────────────────────────────────────

async def _score_chain(client: httpx.AsyncClient, chain: Dict) -> Optional[float]:
    """
    POST the chain to the llm-guard-kit /score endpoint.
    Returns risk float or None on failure (guard errors must not block responses).
    """
    guard_url = _config.get("guard_url", DEFAULT_GUARD_URL)
    try:
        resp = await client.post(
            f"{guard_url}/score",
            json=chain,
            timeout=3.0,
        )
        if resp.status_code == 200:
            data = resp.json()
            return float(data.get("risk", data.get("score", 0.0)))
    except Exception:
        pass
    return None


# ── Proxy helpers ─────────────────────────────────────────────────────────────

def _upstream_headers(request: Request) -> Dict[str, str]:
    """Forward Authorization header from the incoming request to upstream."""
    headers: Dict[str, str] = {"Content-Type": "application/json"}
    auth = request.headers.get("Authorization")
    if auth:
        headers["Authorization"] = auth
    else:
        key = _config.get("upstream_api_key") or os.environ.get("UPSTREAM_API_KEY", "")
        if key:
            headers["Authorization"] = f"Bearer {key}"
    return headers


def _extract_question(body: Dict) -> str:
    """Best-effort question extraction from chat messages."""
    messages = body.get("messages", [])
    # Last user message is the most recent question
    for msg in reversed(messages):
        if msg.get("role") == "user":
            content = msg.get("content", "")
            if isinstance(content, str):
                return content[:500]
            if isinstance(content, list):
                # multimodal content blocks
                for block in content:
                    if isinstance(block, dict) and block.get("type") == "text":
                        return block.get("text", "")[:500]
    return messages[-1].get("content", "")[:500] if messages else ""


# ── Routes ────────────────────────────────────────────────────────────────────

@app.post("/v1/chat/completions")
async def chat_completions(request: Request):
    body         = await request.json()
    upstream_url = _config.get("upstream_base_url", DEFAULT_UPSTREAM)
    headers      = _upstream_headers(request)

    async with httpx.AsyncClient() as client:
        # Forward to upstream
        try:
            upstream_resp = await client.post(
                f"{upstream_url}/v1/chat/completions",
                json=body,
                headers=headers,
                timeout=60.0,
            )
        except httpx.RequestError as exc:
            raise HTTPException(status_code=502, detail=f"Upstream error: {exc}")

        if upstream_resp.status_code != 200:
            return Response(
                content=upstream_resp.content,
                status_code=upstream_resp.status_code,
                media_type="application/json",
            )

        data = upstream_resp.json()

        # Extract the assistant message
        choices = data.get("choices", [])
        if not choices:
            return JSONResponse(data)

        message      = choices[0].get("message", {})
        final_answer = message.get("content") or ""

        # Reconstruct chain + score
        question = _extract_question(body)
        chain    = _build_chain(question, message, final_answer)
        risk     = await _score_chain(client, chain)

        threshold = float(_config.get("guard_threshold", DEFAULT_THRESHOLD))

        # Inject guard metadata into the response (non-breaking extension)
        if risk is not None:
            data["choices"][0]["message"]["guard"] = {
                "risk":    round(risk, 4),
                "alert":   risk > threshold,
                "version": "llm-guard-kit",
            }

        risk_str  = f"{risk:.4f}" if risk is not None else "unavailable"
        alert_str = str(risk > threshold).lower() if risk is not None else "false"

        return JSONResponse(
            data,
            headers={
                "X-LLMGuard-Risk":  risk_str,
                "X-LLMGuard-Alert": alert_str,
            },
        )


@app.post("/v1/completions")
async def completions(request: Request):
    """Legacy completions endpoint — proxied without chain reconstruction."""
    body         = await request.json()
    upstream_url = _config.get("upstream_base_url", DEFAULT_UPSTREAM)
    headers      = _upstream_headers(request)

    async with httpx.AsyncClient() as client:
        try:
            upstream_resp = await client.post(
                f"{upstream_url}/v1/completions",
                json=body,
                headers=headers,
                timeout=60.0,
            )
        except httpx.RequestError as exc:
            raise HTTPException(status_code=502, detail=f"Upstream error: {exc}")

        return Response(
            content=upstream_resp.content,
            status_code=upstream_resp.status_code,
            media_type="application/json",
        )


@app.get("/v1/models")
async def list_models(request: Request):
    """Forward models list from upstream."""
    upstream_url = _config.get("upstream_base_url", DEFAULT_UPSTREAM)
    headers      = _upstream_headers(request)

    async with httpx.AsyncClient() as client:
        try:
            resp = await client.get(
                f"{upstream_url}/v1/models",
                headers=headers,
                timeout=10.0,
            )
            return Response(content=resp.content, status_code=resp.status_code,
                            media_type="application/json")
        except httpx.RequestError as exc:
            raise HTTPException(status_code=502, detail=f"Upstream error: {exc}")


@app.get("/health")
async def health():
    return {
        "status":       "ok",
        "upstream":     _config.get("upstream_base_url", DEFAULT_UPSTREAM),
        "guard_url":    _config.get("guard_url", DEFAULT_GUARD_URL),
        "threshold":    _config.get("guard_threshold", DEFAULT_THRESHOLD),
    }


# ── CLI entry point ───────────────────────────────────────────────────────────

def run():
    import uvicorn

    parser = argparse.ArgumentParser(description="llm-guard-kit OpenAI proxy")
    parser.add_argument(
        "--upstream",
        default=os.environ.get("UPSTREAM_BASE_URL", DEFAULT_UPSTREAM),
        help="Upstream OpenAI-compatible base URL",
    )
    parser.add_argument(
        "--upstream-key",
        default=os.environ.get("UPSTREAM_API_KEY", ""),
        help="API key forwarded to upstream",
    )
    parser.add_argument(
        "--guard-url",
        default=os.environ.get("GUARD_URL", DEFAULT_GUARD_URL),
        help="llm-guard-kit HTTP API URL",
    )
    parser.add_argument(
        "--threshold",
        type=float,
        default=float(os.environ.get("GUARD_THRESHOLD", DEFAULT_THRESHOLD)),
        help="Risk threshold for X-LLMGuard-Alert (default: 0.65)",
    )
    parser.add_argument(
        "--port",
        type=int,
        default=int(os.environ.get("PROXY_PORT", DEFAULT_PORT)),
        help="Port to listen on (default: 8001)",
    )
    parser.add_argument("--host", default="0.0.0.0")
    args = parser.parse_args()

    _config["upstream_base_url"] = args.upstream.rstrip("/")
    _config["upstream_api_key"]  = args.upstream_key
    _config["guard_url"]         = args.guard_url.rstrip("/")
    _config["guard_threshold"]   = args.threshold

    print(f"llm-guard-kit OpenAI proxy")
    print(f"  Upstream:  {_config['upstream_base_url']}")
    print(f"  Guard API: {_config['guard_url']}")
    print(f"  Threshold: {_config['guard_threshold']}")
    print(f"  Listening: http://{args.host}:{args.port}")

    uvicorn.run(app, host=args.host, port=args.port)


if __name__ == "__main__":
    run()
