import os
from typing import AsyncGenerator

import httpx
from fastapi import Depends, FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse

from app.auth import get_discovery, settings, validate_token
from app.schemas import AuthCallbackRequest, AuthConfigResponse, ChatRequest

OLLAMA_API_KEY = os.getenv("OLLAMA_API_KEY")
OLLAMA_API_URL = os.getenv("OLLAMA_API_URL", "https://api.ollama.com/v1/chat/completions")
if not OLLAMA_API_KEY:
    raise RuntimeError("Environment variable OLLAMA_API_KEY is required")

app = FastAPI(
    title="OIDC-Protected Ollama Proxy",
    description="FastAPI backend validating OIDC JWTs and streaming Ollama Cloud responses over SSE.",
    version="0.1.0",
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

async def event_stream(request: ChatRequest) -> AsyncGenerator[str, None]:
    payload = {
        "model": request.model,
        "messages": [message.dict() for message in request.messages],
        "stream": True,
    }
    headers = {
        "Authorization": f"Bearer {OLLAMA_API_KEY}",
        "Content-Type": "application/json",
        "Accept": "text/event-stream",
    }

    async with httpx.AsyncClient(timeout=None) as client:
        async with client.stream("POST", OLLAMA_API_URL, json=payload, headers=headers) as response:
            if response.status_code >= 400:
                text = await response.aread()
                raise HTTPException(status_code=502, detail=f"Ollama Cloud error: {text.decode('utf-8', 'ignore')}" )

            async for chunk in response.aiter_text():
                if not chunk:
                    continue
                for line in chunk.splitlines():
                    if line.strip():
                        yield f"data: {line}\n\n"

@app.get("/health")
async def health() -> dict:
    return {"status": "ok"}

@app.get("/auth/config", response_model=AuthConfigResponse)
async def auth_config() -> AuthConfigResponse:
    discovery = get_discovery()
    return AuthConfigResponse(
        issuer=discovery["issuer"],
        authorization_endpoint=discovery["authorization_endpoint"],
        token_endpoint=discovery["token_endpoint"],
    )

@app.post("/auth/callback")
async def auth_callback(payload: AuthCallbackRequest):
    discovery = get_discovery()
    token_endpoint = discovery.get("token_endpoint")
    if not token_endpoint:
        raise HTTPException(status_code=502, detail="OIDC discovery document missing token_endpoint")

    data = {
        "grant_type": "authorization_code",
        "code": payload.code,
        "code_verifier": payload.code_verifier,
        "redirect_uri": str(payload.redirect_uri),
        "client_id": settings.oidc_client_id,
    }

    # Some OIDC providers expect client credentials via HTTP Basic auth
    auth = None
    if settings.oidc_client_secret:
        auth = (settings.oidc_client_id, settings.oidc_client_secret)

    async with httpx.AsyncClient(timeout=15.0) as client:
        response = await client.post(token_endpoint, data=data, headers={"Accept": "application/json"}, auth=auth)

        # Some providers only accept client authentication via POST body
        if response.status_code != 200 and auth is not None:
            # try falling back to client_secret in the request body
            data_with_secret = dict(data)
            data_with_secret["client_secret"] = settings.oidc_client_secret
            response = await client.post(token_endpoint, data=data_with_secret, headers={"Accept": "application/json"})

    if response.status_code != 200:
        raise HTTPException(status_code=400, detail=f"OIDC token exchange failed: {response.text}")

    return response.json()

@app.post("/chat", response_class=StreamingResponse)
async def chat(request: ChatRequest, token=Depends(validate_token)):
    return StreamingResponse(event_stream(request), media_type="text/event-stream")