from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
import httpx
import os

app = FastAPI(title="Ollama with MCP Tools")

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

# Ollama API configuration
OLLAMA_API_URL = "http://localhost:11434/api"

# MCP Tools integration
class MCPTools:
    @staticmethod
    async def process_with_mcp(prompt: str):
        # Example MCP tool integration
        # Replace with actual MCP tool implementation
        return {"status": "processed_with_mcp", "result": f"Processed: {prompt}"}

# Ollama API client
class OllamaClient:
    @staticmethod
    async def generate(prompt: str, model: str = "gemma3n:2b"):
        async with httpx.AsyncClient() as client:
            try:
                response = await client.post(
                    f"{OLLAMA_API_URL}/generate",
                    json={
                        "model": model,
                        "prompt": prompt,
                        "stream": False
                    },
                    timeout=60.0
                )
                return response.json()
            except Exception as e:
                raise HTTPException(status_code=500, detail=str(e))

# API Endpoints
@app.get("/")
async def read_root():
    return {"status": "Ollama API with MCP Tools is running"}

@app.post("/api/generate")
async def generate_text(prompt: str, use_mcp: bool = False):
    if use_mcp:
        mcp_result = await MCPTools.process_with_mcp(prompt)
        return mcp_result
    
    ollama_response = await OllamaClient().generate(prompt)
    return {"response": ollama_response}

# Health check endpoint
@app.get("/health")
async def health_check():
    try:
        async with httpx.AsyncClient() as client:
            response = await client.get(f"{OLLAMA_API_URL}/tags")
            return {"status": "healthy", "ollama_status": "running" if response.status_code == 200 else "error"}
    except Exception as e:
        return {"status": "error", "message": str(e)}

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