import os
from fastapi import FastAPI, Request
import uvicorn
from litellm import acompletion

# Strip the API key to kill any hidden spaces or newlines!
os.environ["GEMINI_API_KEY"] = os.getenv("GEMINI_API_KEY", "").strip()

app = FastAPI()

# --- MEMORY BANK ---
user_engines = {}

async def run_agent(prompt: str, model_name: str) -> str:
    try:
        response = await acompletion(
            model=model_name, 
            messages=[{"role": "user", "content": prompt}],
            temperature=0.2
        )
        return response.choices[0].message.content or "No response."
    except Exception as e:
        return f"⚠️ Engine Error: {str(e)}"

@app.api_route("/webhook", methods=["GET", "POST", "PUT"])
async def telegram_webhook(request: Request):
    if request.method == "GET":
        return {"status": "alive"}
        
    try:
        data = await request.json()
        
        if "message" in data and "text" in data["message"]:
            chat_id = data["message"]["chat"]["id"]
            user_msg = data["message"]["text"].strip()
            
            # 1. GEAR SHIFTER
            # Swapped to the active 2026 models
            if user_msg.lower() == "/pro":
                user_engines[chat_id] = "gemini/gemini-2.0-pro" 
                return {"method": "sendMessage", "chat_id": chat_id, "text": "🧠 Engine Switched: Gemini 2.0 Pro."}
            elif user_msg.lower() == "/flash":
                user_engines[chat_id] = "gemini/gemini-2.5-flash"
                return {"method": "sendMessage", "chat_id": chat_id, "text": "⚡ Engine Switched: Gemini 2.5 Flash."}
            
            # 2. RUN AI
            current_model = user_engines.get(chat_id, "gemini/gemini-2.5-flash") 
            response_text = await run_agent(user_msg, current_model)
            
            # 3. DIRECT WEBHOOK REPLY
            return {
                "method": "sendMessage",
                "chat_id": chat_id,
                "text": response_text
            }
            
        return {"status": "ok"}
    except Exception as e:
        print(f"CRASH: {repr(e)}")
        return {"status": "error"}

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