# blaxel_main.py (Updated: New Provider Keys + Core Logic Only)
from fastapi import FastAPI
from pydantic import BaseModel
from sse_starlette.sse import EventSourceResponse
import os
import uvicorn
from typing import Dict, Optional

from agent_logic import StrategicSelectorAgent

app = FastAPI()

# 1. Update Pydantic model to accept all provider keys
class ApiKeys(BaseModel):
    google: str
    anthropic: Optional[str] = None
    sambanova: Optional[str] = None
    # New providers added:
    openai: Optional[str] = None
    nebius: Optional[str] = None

class ProblemRequest(BaseModel):
    problem: str
    # Expansion flag removed - we are sticking to the robust core loop
    keys: ApiKeys

@app.post("/solve_problem")
async def solve_problem(request: ProblemRequest):
    """
    Standard endpoint. Initializes the agent with all provided keys
    and streams the solution process.
    """
    async def stream_solution():
        try:
            # Pass the dictionary of keys (including new ones) to the agent
            # request.keys.dict() will include openai and nebius if sent
            mudabbir_ai = StrategicSelectorAgent(api_keys=request.keys.dict())
            
            async for status_update in mudabbir_ai.solve(request.problem):
                yield status_update
        except Exception as e:
            yield f"AGENT ERROR: {e}"
            print(f"AGENT ERROR: {e}") 

    return EventSourceResponse(stream_solution())

if __name__ == "__main__":
    HOST = os.getenv("BL_SERVER_HOST", "0.0.0.0")
    PORT = int(os.getenv("BL_SERVER_PORT", "8080"))
    uvicorn.run(app, host=HOST, port=PORT)