from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import StreamingResponse
import httpx
import logging

app = FastAPI()
OLLAMA_BASE_URL = "http://localhost:11434"
client = httpx.AsyncClient(base_url=OLLAMA_BASE_URL)

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("uvicorn")

@app.middleware("http")
async def log_requests(request: Request, call_next):
    logger.info(f"Request: {request.method} {request.url}")
    response = await call_next(request)
    return response

@app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE"])
async def reverse_proxy(request: Request, path: str):
    try:
        # Build the Ollama URL
        url = f"{OLLAMA_BASE_URL}/{path}"
        
        # Forward headers excluding those that might cause issues
        headers = {
            key: value
            for key, value in request.headers.items()
            if key.lower() not in ["host", "content-length"]
        }
        
        # Forward the request to Ollama
        req = client.build_request(
            method=request.method,
            url=url,
            headers=headers,
            content=await request.body(),
            params=request.query_params
        )
        
        response = await client.send(req, stream=True)
        
        # Handle streaming responses
        if "text/event-stream" in response.headers.get("content-type", ""):
            return StreamingResponse(
                response.aiter_bytes(),
                media_type=response.headers.get("content-type"),
                headers=dict(response.headers)
            )
            
        return StreamingResponse(
            response.aiter_bytes(),
            media_type=response.headers.get("content-type"),
            headers=dict(response.headers)
        )
        
    except httpx.ConnectError:
        raise HTTPException(
            status_code=503,
            detail="Ollama server unavailable"
        )
    except Exception as e:
        logger.error(f"Proxy error: {str(e)}")
        raise HTTPException(
            status_code=500,
            detail=f"Proxy error: {str(e)}"
        )

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)