# blaxel_main.py (New Secure Version)
from fastapi import FastAPI
from pydantic import BaseModel
from sse_starlette.sse import EventSourceResponse
import os
import uvicorn
from typing import Dict, Optional

# 1. Import our "Brain"
from agent_logic import StrategicSelectorAgent

app = FastAPI()

# 2. Update Pydantic model to accept keys (some are optional)
class ApiKeys(BaseModel):
    google: str
    anthropic: Optional[str] = None
    sambanova: Optional[str] = None

class ProblemRequest(BaseModel):
    problem: str
    keys: ApiKeys

@app.post("/solve_problem")
async def solve_problem(request: ProblemRequest):
    """
    This endpoint now creates a *new* agent for every request,
    and passes the user's API keys to it.
    This is stateless and secure.
    """
    
    async def stream_solution():
        try:
            # 3. Pass keys to the agent's constructor
            mudabbir_ai = StrategicSelectorAgent(api_keys=request.keys.model_dump())
            
            async for status_update in mudabbir_ai.solve(request.problem):
                yield status_update
        except Exception as e:
            # This will catch errors (like invalid keys)
            yield f"AGENT ERROR: {e}"
            print(f"AGENT ERROR: {e}") # Also log to backend

    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)