import logging
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://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",
)

logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)

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, follow_redirects=True) 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": payload.redirect_uri,
    }

    # 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)
    else:
        data["client_id"] = settings.oidc_client_id

    async with httpx.AsyncClient(timeout=15.0) as client:
        logger.debug("Token endpoint: %s", token_endpoint)
        logger.debug("Token request data: %s", data)
        logger.debug("Using auth: %s", "basic" if auth is not None else "none")

        response = await client.post(token_endpoint, data=data, headers={"Accept": "application/json"}, auth=auth)
        logger.debug("Token response status: %s", response.status_code)
        logger.debug("Token response headers: %s", response.headers)
        logger.debug("Token response text: %s", response.text)

        # Some providers only accept client authentication via POST body
        if response.status_code != 200 and auth is not None:
            data_with_secret = dict(data)
            data_with_secret["client_id"] = settings.oidc_client_id
            data_with_secret["client_secret"] = settings.oidc_client_secret
            logger.debug("Retrying with client_secret_post: %s", data_with_secret)
            response = await client.post(token_endpoint, data=data_with_secret, headers={"Accept": "application/json"})
            logger.debug("Token response status (fallback): %s", response.status_code)
            logger.debug("Token response headers (fallback): %s", response.headers)
            logger.debug("Token response text (fallback): %s", response.text)

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

    result = response.json()
    logger.debug("OIDC callback result keys: %s", list(result.keys()))
    if "access_token" not in result:
        logger.warning("OIDC token response missing access_token; may return only id_token or an opaque token")
    return result

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