from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import List, Dict, Optional
import asyncio
import aiohttp
import json
import time
import logging
from fastapi.responses import StreamingResponse
import uvicorn

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

# Configuration
POLLINATIONS_API_URL = "https://text.pollinations.ai/openai"
API_KEY = "wPGHlU-7pPYlOetQ"
MAX_CONTEXT_MESSAGES = 15

# FastAPI app
app = FastAPI(
    title="AI Assistant API",
    description="Server API for AI Assistant powered by Pollinations",
    version="1.0.0"
)

# CORS middleware - आपके frontend के लिए जरूरी
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Production में specific domains add करें
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Request models
class ChatMessage(BaseModel):
    role: str
    content: str

class ChatRequest(BaseModel):
    type: str
    prompt: str
    model: Optional[str] = "openai"
    conversation: Optional[List[ChatMessage]] = []
    stream: Optional[bool] = True
    private: Optional[bool] = True
    referrer: Optional[str] = API_KEY

class ChatResponse(BaseModel):
    success: bool
    choices: Optional[List[Dict]] = None
    error: Optional[str] = None
    model: Optional[str] = None
    timestamp: Optional[int] = None

@app.get("/")
async def root():
    """Health check endpoint"""
    return {
        "status": "running",
        "message": "AI Assistant API Server",
        "timestamp": int(time.time())
    }

@app.post("/api/chat")
async def chat_endpoint(request: ChatRequest):
    """Main chat endpoint - exactly like your server.php"""
    
    if not request.prompt or not request.prompt.strip():
        raise HTTPException(status_code=400, detail="Prompt cannot be empty")
    
    try:
        if request.stream:
            return StreamingResponse(
                generate_streaming_response(request),
                media_type="text/event-stream",
                headers={
                    "Cache-Control": "no-cache",
                    "Connection": "keep-alive",
                    "Access-Control-Allow-Origin": "*",
                    "Access-Control-Allow-Methods": "POST, GET, OPTIONS",
                    "Access-Control-Allow-Headers": "Content-Type"
                }
            )
        else:
            response = await generate_non_streaming_response(request)
            return response
            
    except Exception as e:
        logger.error(f"Chat error: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Server error: {str(e)}")

async def generate_streaming_response(request: ChatRequest):
    """Generate streaming response - exactly like your PHP version"""
    
    # System message
    system_message = {
        'role': 'system',
        'content': 'आप एक helpful AI assistant हैं। User को Hindi और English दोनों languages में helpful और accurate answers देते हैं। आप friendly, conversational और natural tone में बात करते हैं। Code माँगने पर proper formatting के साथ दें। हमेशा relevant और उपयोगी जवाब दें।'
    }
    
    messages = [system_message]
    
    # Add conversation history
    for msg in request.conversation[-MAX_CONTEXT_MESSAGES:]:
        if msg.role in ['user', 'assistant'] and msg.content.strip():
            content = msg.content[:2000]  # Limit content length
            messages.append({
                'role': msg.role,
                'content': content
            })
    
    # Add current prompt
    messages.append({
        'role': 'user',
        'content': request.prompt
    })
    
    payload = {
        'model': request.model,
        'messages': messages,
        'temperature': 0.7,
        'max_tokens': 1000,
        'top_p': 0.9,
        'stream': True,
        'private': request.private,
        'referrer': request.referrer,
        'seed': int(time.time()) % 1000000
    }
    
    try:
        async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=60)) as session:
            async with session.post(
                POLLINATIONS_API_URL,
                json=payload,
                headers={
                    'Content-Type': 'application/json',
                    'Accept': 'text/event-stream',
                    'User-Agent': 'AI-Assistant-API/1.0'
                }
            ) as response:
                
                if response.status != 200:
                    error_text = await response.text()
                    yield f"data: {json.dumps({'error': f'API returned HTTP {response.status}: {error_text}'})}\n\n"
                    return
                
                async for line in response.content:
                    line_text = line.decode('utf-8')
                    yield line_text
                
                yield "data: [DONE]\n\n"
                
    except asyncio.TimeoutError:
        yield f"data: {json.dumps({'error': 'Request timeout - कृपया दोबारा कोशिश करें'})}\n\n"
    except Exception as e:
        yield f"data: {json.dumps({'error': f'Network error: {str(e)}'})}\n\n"

async def generate_non_streaming_response(request: ChatRequest):
    """Generate non-streaming response"""
    
    system_message = {
        'role': 'system',
        'content': 'आप एक helpful AI assistant हैं। User को Hindi और English दोनों languages में helpful और accurate answers देते हैं। आप friendly, conversational और natural tone में बात करते हैं। Code माँगने पर proper formatting के साथ दें। हमेशा relevant और उपयोगी जवाब दें।'
    }
    
    messages = [system_message]
    
    for msg in request.conversation[-MAX_CONTEXT_MESSAGES:]:
        if msg.role in ['user', 'assistant'] and msg.content.strip():
            content = msg.content[:2000]
            messages.append({
                'role': msg.role,
                'content': content
            })
    
    messages.append({
        'role': 'user',
        'content': request.prompt
    })
    
    payload = {
        'model': request.model,
        'messages': messages,
        'temperature': 0.7,
        'max_tokens': 1000,
        'top_p': 0.9,
        'stream': False,
        'private': request.private,
        'referrer': request.referrer,
        'seed': int(time.time()) % 1000000
    }
    
    try:
        async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=30)) as session:
            async with session.post(
                POLLINATIONS_API_URL,
                json=payload,
                headers={
                    'Content-Type': 'application/json',
                    'User-Agent': 'AI-Assistant-API/1.0'
                }
            ) as response:
                
                if response.status != 200:
                    error_text = await response.text()
                    return ChatResponse(
                        success=False,
                        error=f"API returned HTTP {response.status}: {error_text}"
                    )
                
                data = await response.json()
                
                return ChatResponse(
                    success=True,
                    choices=data.get('choices', []),
                    model=request.model,
                    timestamp=int(time.time())
                )
                
    except Exception as e:
        logger.error(f"API Error: {str(e)}")
        return ChatResponse(
            success=False,
            error=f"Network error: {str(e)}"
        )

# Health check endpoints
@app.get("/health")
async def health_check():
    return {"status": "healthy", "timestamp": int(time.time())}

@app.get("/api/status")
async def api_status():
    return {
        "api_version": "1.0.0",
        "pollinations_endpoint": POLLINATIONS_API_URL,
        "max_context_messages": MAX_CONTEXT_MESSAGES,
        "supported_models": ["openai"],
        "timestamp": int(time.time())
    }

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