import sys
from fastapi import FastAPI, Request
from fastapi.responses import HTMLResponse, StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from curl_cffi import requests as cffi_requests

app = FastAPI()

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

TARGET_URL = "https://agentrouter.org/v1/chat/completions"

@app.post("/v1/chat/completions")
async def proxy_chat(request: Request):
    body = await request.body()
    
    headers = {
        "Authorization": request.headers.get("Authorization", ""),
        "Content-Type": "application/json",
        "Accept": "text/event-stream",
        "Originator": "codex_cli_rs",
        "User-Agent": "codex_cli_rs/0.101.0 (Mac OS 26.0.1; arm64) Apple_Terminal/464",
        "Version": "0.101.0"
    }

    async def stream_generator():
        try:
            async with cffi_requests.AsyncSession(impersonate="chrome110") as session:
                response = await session.post(
                    TARGET_URL, 
                    headers=headers, 
                    data=body, 
                    stream=True, 
                    timeout=180.0
                )
                
                # Agar saaf error 400 ya 401 aaye
                if response.status_code != 200:
                    error_content = b""
                    async for chunk in response.aiter_content():
                        if chunk:
                            error_content += chunk
                    safe_error = error_content.decode('utf-8', 'ignore').replace('"', "'").replace('\n', ' ')
                    error_msg = f"🚨 API Status Error ({response.status_code}): {safe_error[:150]}..."
                    yield f'data: {{"choices":[{{"delta":{{"content":"{error_msg}"}}}}]}}\n\n'.encode('utf-8')
                    yield b'data: [DONE]\n\n'
                    return
                
                # Agar 200 OK aaye, toh Pehla Chunk check karein (Spy Camera)
                first_chunk = True
                async for chunk in response.aiter_content():
                    if chunk:
                        if first_chunk:
                            # 1. Logs mein exact AgentRouter ka jawab print karega
                            print(f"\n--- 🔎 AGENTROUTER RESPONSE LOG ---\n{chunk[:500]}\n-----------------------------------\n", file=sys.stderr, flush=True)
                            first_chunk = False
                            
                            # 2. Agar API Key ya Quota ka chupa hua error hai (Jo JSON form mein aata hai)
                            if chunk.lstrip().startswith(b'{') and b'"error"' in chunk:
                                safe_err = chunk.decode('utf-8', 'ignore').replace('"', "'").replace('\n', ' ')
                                yield f'data: {{"choices":[{{"delta":{{"content":"🚨 **API/Key Error:** {safe_err}"}}}}]}}\n\n'.encode('utf-8')
                                yield b'data: [DONE]\n\n'
                                return
                                
                        # Normal AI Text Data
                        yield chunk
                        
        except Exception as e:
            error_msg = f"🚨 Network Drop: {str(e)}"
            yield f'data: {{"choices":[{{"delta":{{"content":"{error_msg}"}}}}]}}\n\n'.encode('utf-8')
            yield b'data: [DONE]\n\n'

    return StreamingResponse(stream_generator(), media_type="text/event-stream")

@app.get("/", response_class=HTMLResponse)
async def read_root():
    with open("index.html", "r", encoding="utf-8") as f:
        return f.read()