from fastapi import FastAPI, Request, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse, JSONResponse, FileResponse
from pydantic import BaseModel
from typing import List, Optional, AsyncGenerator
import json
import time
import uuid
import logging
import asyncio
import os

# استيراد g4f
import g4f
from g4f import ChatCompletion

# =====================================================
# LOGGING
# =====================================================
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# =====================================================
# CONFIG
# =====================================================
API_KEY = os.environ.get("API_KEY", "sk-your-secret-key") # يدعم متغيرات البيئة في HF

# =====================================================
# FASTAPI
# =====================================================
app = FastAPI(
    title="9Router AI Gateway",
    version="6.0.0"
)

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

# =====================================================
# MODELS
# =====================================================
class Message(BaseModel):
    role: str
    content: str

class ChatRequest(BaseModel):
    model: str = "gpt-3.5-turbo"
    messages: List[Message]
    stream: bool = False
    temperature: Optional[float] = 0.7
    max_tokens: Optional[int] = 4096

# =====================================================
# AUTH (متوافق مع الواجهة)
# =====================================================
def verify_api_key(req: Request):
    auth = req.headers.get("Authorization")
    cf_auth = req.headers.get("cf-aig-authorization") # دعم Cloudflare Tunnel من الواجهة
    
    # إذا لم يتم إرسال أي مفتاح، نسمح بالمرور (للاختبار المحلي)
    if not auth and not cf_auth:
        return True
    
    token = ""
    if auth and auth.startswith("Bearer "):
        token = auth.replace("Bearer ", "").strip()
    elif cf_auth and cf_auth.startswith("Bearer "):
        token = cf_auth.replace("Bearer ", "").strip()
        
    if token and token != API_KEY:
        raise HTTPException(status_code=403, detail="Invalid API Key")
    
    return True

# =====================================================
# SERVE UI (تقديم واجهة 9Router)
# =====================================================
@app.get("/")
async def serve_ui():
    # هذا سيقوم بعرض ملف index.html تلقائياً عند فتح رابط الـ Space
    return FileResponse("index.html")

# =====================================================
# API MODELS
# =====================================================
@app.get("/v1/models")
async def get_models():
    models_data = [
        {"id": "gpt-3.5-turbo", "object": "model", "owned_by": "g4f"},
        {"id": "gpt-4", "object": "model", "owned_by": "g4f"},
        {"id": "gpt-4o-mini", "object": "model", "owned_by": "g4f"},
        {"id": "gpt-4o", "object": "model", "owned_by": "g4f"},
        {"id": "claude-3-haiku", "object": "model", "owned_by": "g4f"},
        {"id": "gemini-pro", "object": "model", "owned_by": "g4f"},
    ]
    return {"object": "list", "data": models_data}

# =====================================================
# CHAT COMPLETIONS
# =====================================================
@app.post("/v1/chat/completions")
async def chat_completions(req: Request, body: ChatRequest):
    verify_api_key(req)
    
    messages = [{"role": m.role, "content": m.content} for m in body.messages]
    
    logger.info(f"Request: model={body.model}, stream={body.stream}")
    
    # =================================================
    # STREAMING
    # =================================================
    if body.stream:
        async def generate_stream() -> AsyncGenerator[str, None]:
            chunk_id = f"chatcmpl-{uuid.uuid4().hex[:8]}"
            
            try:
                # إضافة Timeout وإزالة المزود الثابت ليختار g4f الأفضل
                response = await asyncio.wait_for(
                    ChatCompletion.create_async(
                        model=body.model,
                        messages=messages,
                        stream=True,
                    ),
                    timeout=60.0
                )
                
                async for chunk in response:
                    if chunk:
                        # إجبار التحويل إلى نص لمنع انهيار JSON
                        content = str(chunk) if not isinstance(chunk, str) else chunk
                        
                        payload = {
                            "id": chunk_id,
                            "object": "chat.completion.chunk",
                            "created": int(time.time()),
                            "model": body.model,
                            "choices": [{
                                "index": 0,
                                "delta": {"content": content},
                                "finish_reason": None
                            }]
                        }
                        yield f"data: {json.dumps(payload, ensure_ascii=False)}\n\n"
                
                # إشارة النهاية
                final_payload = {
                    "id": chunk_id,
                    "object": "chat.completion.chunk",
                    "created": int(time.time()),
                    "model": body.model,
                    "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]
                }
                yield f"data: {json.dumps(final_payload)}\n\n"
                yield "data: [DONE]\n\n"
                
            except asyncio.TimeoutError:
                logger.error("Streaming timeout")
                error_msg = {"error": {"message": "Provider timeout", "type": "timeout"}}
                yield f"data: {json.dumps(error_msg)}\n\n"
                yield "data: [DONE]\n\n"
            except Exception as e:
                logger.error(f"Streaming error: {e}", exc_info=True)
                error_msg = {"error": {"message": str(e), "type": "stream_error"}}
                yield f"data: {json.dumps(error_msg)}\n\n"
                yield "data: [DONE]\n\n"
        
        return StreamingResponse(
            generate_stream(),
            media_type="text/event-stream",
            headers={
                "Cache-Control": "no-cache",
                "Connection": "keep-alive",
            }
        )
    
    # =================================================
    # NORMAL RESPONSE
    # =================================================
    try:
        response = await asyncio.wait_for(
            ChatCompletion.create_async(
                model=body.model,
                messages=messages,
                stream=False,
            ),
            timeout=90.0
        )
        
        # إجبار التحول لنص
        response_text = str(response) if not isinstance(response, str) else response

        return JSONResponse({
            "id": f"chatcmpl-{uuid.uuid4().hex[:8]}",
            "object": "chat.completion",
            "created": int(time.time()),
            "model": body.model,
            "choices": [{
                "index": 0,
                "message": {
                    "role": "assistant",
                    "content": response_text
                },
                "finish_reason": "stop"
            }],
            "usage": {
                "prompt_tokens": 0,
                "completion_tokens": 0,
                "total_tokens": 0
            }
        })
        
    except asyncio.TimeoutError:
        raise HTTPException(status_code=504, detail="Provider timed out")
    except Exception as e:
        logger.error(f"Chat error: {e}", exc_info=True)
        raise HTTPException(status_code=500, detail=f"Provider error: {str(e)}")

# =====================================================
# RUN
# =====================================================
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(
        app,
        host="0.0.0.0",
        port=7860, # المنفذ المطلوب لـ Hugging Face
        log_level="info"
    )