app.py

import os
import re
import logging
import uuid
import time
from datetime import datetime, timezone, timedelta
from collections import defaultdict
from typing import Optional, Dict, Any
import asyncio
from concurrent.futures import ThreadPoolExecutor

from fastapi import FastAPI, HTTPException, Body, BackgroundTasks, Path, Request
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, Field

import openai # For your custom API
import google.generativeai as genai # For Gemini API
from google.generativeai.types import GenerationConfig

# --- Logging Configuration ---
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger(__name__)

# --- Configuration ---
CUSTOM_API_BASE_URL_DEFAULT = "https://api-q3ieh5raqfuad9o8.aistudio-app.com/v1"
CUSTOM_API_MODEL_DEFAULT = "gemma3:27b"
DEFAULT_GEMINI_MODEL = "gemini-2.0-flash"
GEMINI_REQUEST_TIMEOUT_SECONDS = 300

# --- In-Memory Task Storage ---
tasks_db: Dict[str, Dict[str, Any]] = {}

# --- Pydantic Models ---
class ChatPayload(BaseModel):
    message: str
    temperature: float = Field(0.6, ge=0.0, le=1.0)

class GeminiTaskRequest(BaseModel):
    message: str
    url: Optional[str] = None
    gemini_model: Optional[str] = None
    api_key: Optional[str] = Field(None, description="Gemini API Key (optional; uses Space secret if not provided)")

class TaskSubmissionResponse(BaseModel):
    task_id: str
    status: str
    task_detail_url: str

class TaskStatusResponse(BaseModel):
    task_id: str
    status: str
    submitted_at: datetime
    last_updated_at: datetime
    result: Optional[str] = None
    error: Optional[str] = None
    # request_params: Optional[Dict[str, Any]] = None # Optionally return original params


# Rate limiting dictionary
class RateLimiter:
    def __init__(self, max_requests: int, time_window: timedelta):
        self.max_requests = max_requests
        self.time_window = time_window
        self.requests: Dict[str, list] = defaultdict(list)
    
    def _cleanup_old_requests(self, user_ip: str) -> None:
        """Remove requests that are outside the time window."""
        current_time = time.time()
        self.requests[user_ip] = [
            timestamp for timestamp in self.requests[user_ip]
            if current_time - timestamp < self.time_window.total_seconds()
        ]
    
    def is_rate_limited(self, user_ip: str) -> bool:
        """Check if the user has exceeded their rate limit."""
        self._cleanup_old_requests(user_ip)
        
        # Get current count after cleanup
        current_count = len(self.requests[user_ip])
        
        # Add current request timestamp (incrementing the count)
        current_time = time.time()
        self.requests[user_ip].append(current_time)
        
        # Check if user has exceeded the maximum requests
        return (current_count + 1) > self.max_requests
    
    def get_current_count(self, user_ip: str) -> int:
        """Get the current request count for an IP."""
        self._cleanup_old_requests(user_ip)
        return len(self.requests[user_ip])


# Initialize rate limiter with 100 requests per day
rate_limiter = RateLimiter(
    max_requests=50,
    time_window=timedelta(days=1)
)

def get_user_ip(request: Request) -> str:
    """Helper function to get user's IP address."""
    forwarded = request.headers.get("X-Forwarded-For")
    if forwarded:
        return forwarded.split(",")[0]
    return request.client.host


class ApiRotator:
    def __init__(self, apis):
        self.apis = apis
        self.last_successful_index = None

    def get_prioritized_apis(self):
        if self.last_successful_index is not None:
            # Move the last successful API to the front
            rotated_apis = (
                [self.apis[self.last_successful_index]] + 
                self.apis[:self.last_successful_index] + 
                self.apis[self.last_successful_index+1:]
            )
            return rotated_apis
        return self.apis

    def update_last_successful(self, index):
        self.last_successful_index = index


# --- FastAPI App Initialization ---
app = FastAPI(
    title="Dual Chat & Async Gemini API",
    description="Made by Cody from chrunos.com.",
    version="2.0.0"
)

# --- Helper Functions ---
def is_video_url_for_gemini(url: Optional[str]) -> bool:
    if not url:
        return False
    # Use raw strings (r"...") for regular expressions to avoid SyntaxWarnings
    youtube_regex = (
        r'(https_?://)?(www\.)?'
        r'(youtube|youtu|youtube-nocookie)\.(com|be)/'  # Changed to raw string
        r'(watch\?v=|embed/|v/|.+\?v=)?([^&=%\?]{11})'  # Changed to raw string
    )
    # This regex was likely fine as it didn't have ambiguous escapes, but good practice to make it raw too
    googleusercontent_youtube_regex = r'https_?://googleusercontent\.com/youtube\.com/\w+'
    
    return re.match(youtube_regex, url) is not None or \
           re.match(googleusercontent_youtube_regex, url) is not None

async def process_gemini_request_background(
    task_id: str, 
    user_message: str, 
    input_url: Optional[str], 
    requested_gemini_model: str, 
    gemini_key_to_use: str
):
    logger.info(f"[Task {task_id}] Starting background Gemini processing. Model: {requested_gemini_model}, URL: {input_url}")
    tasks_db[task_id]["status"] = "PROCESSING"
    tasks_db[task_id]["last_updated_at"] = datetime.now(timezone.utc)

    try:
        genai.configure(api_key=gemini_key_to_use)
        
        model_instance = genai.GenerativeModel(model_name=requested_gemini_model)
        
        content_parts = [{"text": user_message}]
        if input_url and is_video_url_for_gemini(input_url):
            logger.info(f"[Task {task_id}] Adding video URL to Gemini content: {input_url}")
            content_parts.append({
                "file_data": {
                    "mime_type": "video/youtube", # Or let Gemini infer
                    "file_uri": input_url
                }
            })
        
        gemini_contents = [{"parts": content_parts}]
        
        generation_config = GenerationConfig(candidate_count=1)
        request_options = {"timeout": GEMINI_REQUEST_TIMEOUT_SECONDS}

        logger.info(f"[Task {task_id}] Sending request to Gemini API...")
        response = await model_instance.generate_content_async(
            gemini_contents,
            stream=False, # Collect full response for async task
            generation_config=generation_config,
            request_options=request_options
        )
        
        # Assuming response.text contains the full aggregated text
        # If using a model version that streams even for non-stream call, aggregate it:
        full_response_text = ""
        if hasattr(response, 'text') and response.text:
            full_response_text = response.text
        elif hasattr(response, 'parts'): # Check for newer API structures if .text is not primary
            for part in response.parts:
                if hasattr(part, 'text'):
                    full_response_text += part.text
        else: # Fallback for safety if structure is unexpected or if it's an iterable of chunks
            # This part might need adjustment based on actual non-streaming response object
            # For now, assuming generate_content_async with stream=False gives a response with .text
            # or we need to iterate if it's still a stream internally for some models
            logger.warning(f"[Task {task_id}] Gemini response structure not as expected or empty. Response: {response}")


        if not full_response_text and response.prompt_feedback and response.prompt_feedback.block_reason:
            block_reason_name = response.prompt_feedback.block_reason.name if hasattr(response.prompt_feedback.block_reason, 'name') else str(response.prompt_feedback.block_reason)
            logger.warning(f"[Task {task_id}] Gemini content blocked: {block_reason_name}")
            tasks_db[task_id]["status"] = "FAILED"
            tasks_db[task_id]["error"] = f"Content blocked by Gemini due to: {block_reason_name}"
        elif full_response_text:
            logger.info(f"[Task {task_id}] Gemini processing successful. Result length: {len(full_response_text)}")
            tasks_db[task_id]["status"] = "COMPLETED"
            tasks_db[task_id]["result"] = full_response_text
        else:
            logger.warning(f"[Task {task_id}] Gemini processing completed but no text content found and no block reason.")
            tasks_db[task_id]["status"] = "FAILED"
            tasks_db[task_id]["error"] = "Gemini returned no content and no specific block reason."

    except Exception as e:
        logger.error(f"[Task {task_id}] Error during Gemini background processing: {e}", exc_info=True)
        tasks_db[task_id]["status"] = "FAILED"
        tasks_db[task_id]["error"] = str(e)
    finally:
        tasks_db[task_id]["last_updated_at"] = datetime.now(timezone.utc)

# --- API Endpoints ---

@app.post("/chat", response_class=StreamingResponse)
async def direct_chat(payload: ChatPayload, request: Request):
    logger.info(f"Direct chat request received. Temperature: {payload.temperature}, Message: '{payload.message[:50]}...'")
    user_ip = get_user_ip(request)
    
    if rate_limiter.is_rate_limited(user_ip):
        current_count = rate_limiter.get_current_count(user_ip)
        raise HTTPException(
            status_code=429,
            detail={
                "error": "You have exceeded the maximum number of requests per day. Please try again tomorrow.",
                "url": "https://t.me/chrunoss"
            }
        )
    custom_api_key_secret = os.getenv("CUSTOM_API_SECRET_KEY")
    custom_api_base_url = os.getenv("CUSTOM_API_BASE_URL", CUSTOM_API_BASE_URL_DEFAULT)
    custom_api_model = os.getenv("CUSTOM_API_MODEL", CUSTOM_API_MODEL_DEFAULT)
    
    if not custom_api_key_secret:
        logger.error("Custom API key ('CUSTOM_API_SECRET_KEY') is not configured for /chat.")
        raise HTTPException(status_code=500, detail="Custom API key not configured.")
    
    async def custom_api_streamer():
        client = None
        try:
            logger.info("Sending request to Custom API for /chat.")
            
            # Use AsyncOpenAI with proper configuration
            from openai import AsyncOpenAI
            client = AsyncOpenAI(
                api_key=custom_api_key_secret,
                base_url=custom_api_base_url,
                timeout=60.0  # Longer timeout for gemma3:27b model
            )
            
            stream = await client.chat.completions.create(
                model=custom_api_model,
                temperature=payload.temperature,
                messages=[{"role": "user", "content": payload.message}],
                stream=True
            )
            
            async for chunk in stream:
                try:
                    # Exact same logic as your working code
                    if hasattr(chunk.choices[0].delta, "reasoning_content") and chunk.choices[0].delta.reasoning_content:
                        yield chunk.choices[0].delta.reasoning_content
                    elif chunk.choices[0].delta.content is not None:  # Handle None explicitly
                        yield chunk.choices[0].delta.content
                        
                except (IndexError, AttributeError) as e:
                    # Skip malformed chunks silently (some APIs send empty chunks)
                    continue
                except Exception as e:
                    logger.warning(f"Skipping chunk due to error: {e}")
                    continue
                    
        except Exception as e:
            logger.error(f"Error during Custom API call for /chat: {e}", exc_info=True)
            
            # Handle specific connection errors with retry suggestion
            if "peer closed connection" in str(e) or "incomplete chunked read" in str(e):
                yield "Connection interrupted. Please try again."
            else:
                yield f"Error processing with Custom API: {str(e)}"
                
        finally:
            if client:
                try:
                    await client.close()
                except Exception as cleanup_error:
                    logger.warning(f"Error closing OpenAI client: {cleanup_error}")
    
    return StreamingResponse(
        custom_api_streamer(),
        media_type="text/plain",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
        }
    )

@app.post("/gemini/submit_task", response_model=TaskSubmissionResponse)
async def submit_gemini_task(request: GeminiTaskRequest, background_tasks: BackgroundTasks):
    task_id = str(uuid.uuid4())
    logger.info(f"Received Gemini task submission. Assigning Task ID: {task_id}. Message: '{request.message[:50]}...'")

    gemini_api_key_from_request = request.api_key
    gemini_api_key_secret = os.getenv("GEMINI_API_KEY")
    key_to_use = gemini_api_key_from_request

    if not key_to_use:
        logger.error(f"[Task {task_id}] Gemini API Key missing for task submission.")
        raise HTTPException(status_code=400, detail="Gemini API Key required.")

    requested_model = request.gemini_model or DEFAULT_GEMINI_MODEL
    
    current_time = datetime.now(timezone.utc)
    tasks_db[task_id] = {
        "status": "PENDING",
        "result": None,
        "error": None,
        "submitted_at": current_time,
        "last_updated_at": current_time,
        "request_params": request.model_dump() # Store original request
    }

    background_tasks.add_task(
        process_gemini_request_background,
        task_id,
        request.message,
        request.url,
        requested_model,
        key_to_use
    )
    
    logger.info(f"[Task {task_id}] Task submitted to background processing.")
    return TaskSubmissionResponse(
        task_id=task_id,
        status="PENDING",
        task_detail_url=f"/gemini/task/{task_id}" # Provide the URL to poll
    )



@app.get("/gemini/task/{task_id}", response_model=TaskStatusResponse)
async def get_gemini_task_status(task_id: str = Path(..., description="The ID of the task to retrieve")):
    logger.info(f"Status query for Task ID: {task_id}")
    task = tasks_db.get(task_id)
    if not task:
        logger.warning(f"Task ID not found: {task_id}")
        raise HTTPException(status_code=404, detail="Task ID not found.")
    
    logger.info(f"[Task {task_id}] Current status: {task['status']}")
    return TaskStatusResponse(
        task_id=task_id,
        status=task["status"],
        submitted_at=task["submitted_at"],
        last_updated_at=task["last_updated_at"],
        result=task.get("result"),
        error=task.get("error"),
        # request_params=task.get("request_params") # Optionally include original params
    )

@app.get("/")
async def read_root():
    logger.info("Root endpoint '/' accessed (health check).")
    return {"message": "API for Direct Chat and Async Gemini Tasks is running."}