routers/openrouter.py

from fastapi import APIRouter, Depends, Query, File, UploadFile, HTTPException
from pydantic import BaseModel
from typing import Annotated
from auth import verify_token
from enum import Enum
import os
import httpx
import base64
from schemas.common import APIResponse, ChatMemoryRequest
from datetime import date, datetime, timedelta # Added imports

router = APIRouter(prefix="/open-router", tags=["open-router"])

open_router_key = os.environ.get('OPEN_ROUTER_KEY', '')
OPENROUTER_API_BASE = "https://openrouter.ai/api/v1"
YOUR_SITE_URL = os.environ.get("YOUR_SITE_URL", "http://localhost") # Default or get from env
YOUR_SITE_NAME = os.environ.get("YOUR_SITE_NAME", "") # Default or get from env

# --- Rate Limiting ---
MAX_CALLS_PER_DAY = 200
call_count = 0
last_reset_date = date.today()

def check_and_increment_rate_limit():
    """Checks if the daily rate limit has been exceeded and increments the count."""
    global call_count, last_reset_date
    today = date.today()

    # Reset counter if it's a new day
    if today != last_reset_date:
        call_count = 0
        last_reset_date = today

    if call_count >= MAX_CALLS_PER_DAY:
        raise HTTPException(status_code=429, detail=f"Daily rate limit of {MAX_CALLS_PER_DAY} calls exceeded.")

    call_count += 1
# --- End Rate Limiting ---


class ModelName(str, Enum):
    deepseek_r1 = "deepseek/deepseek-r1:free"
    gemini_2_flash_lite = "google/gemini-2.0-flash-lite-preview-02-05:free"
    gemini_2_flash = "google/gemini-2.0-flash-exp:free"
    gemini_2_pro = "google/gemini-2.0-pro-exp-02-05:free"
    gemini_2_flash_thinking = "google/gemini-2.0-flash-thinking-exp-1219:free"
    gemini_2_5_pro = "google/gemini-2.5-pro-exp-03-25:free"
    gemma_3_27b = "google/gemma-3-27b-it:free"
    llama_3_3 = "meta-llama/llama-3.3-70b-instruct:free"
    llama_4_scout = "meta-llama/llama-4-scout:free"
    llama_4_maverick = "meta-llama/llama-4-maverick:free"
    mistral_small_3 ="mistralai/mistral-small-24b-instruct-2501:free"
    qwen_qwq_32b = "qwen/qwq-32b:free"

class MultiModelName(str, Enum):
    qwen_vl_plus = "qwen/qwen-vl-plus:free"
    qwen_vl_72b = "qwen/qwen2.5-vl-72b-instruct:free"
    gemini_2_flash_lite = "google/gemini-2.0-flash-lite-preview-02-05:free"
    gemini_2_flash = "google/gemini-2.0-flash-exp:free"
    gemini_2_pro = "google/gemini-2.0-pro-exp-02-05:free"
    gemini_2_flash_thinking = "google/gemini-2.0-flash-thinking-exp-1219:free"
    gemini_2_5_pro = "google/gemini-2.5-pro-exp-03-25:free"
    gemma_3_27b = "google/gemma-3-27b-it:free"
    llama_3_2_vision = "meta-llama/llama-3.2-11b-vision-instruct:free"
    llama_4_scout = "meta-llama/llama-4-scout:free"
    llama_4_maverick = "meta-llama/llama-4-maverick:free"


async def _call_openrouter(payload: dict):
    # Check rate limit before making the call
    try:
        check_and_increment_rate_limit()
    except HTTPException as e:
        raise e # Re-raise the 429 error

    if not open_router_key:
        raise HTTPException(status_code=500, detail="OPEN_ROUTER_KEY not configured")

    async with httpx.AsyncClient() as client:
        try:
            response = await client.post(
                url=f"{OPENROUTER_API_BASE}/chat/completions",
                headers={
                    "Authorization": f"Bearer {open_router_key}",
                    "Content-Type": "application/json",
                    "HTTP-Referer": YOUR_SITE_URL,
                    "X-Title": YOUR_SITE_NAME,
                },
                json=payload,
                timeout=60 # Add a timeout
            )
            response.raise_for_status()
            return response.json()
        except httpx.HTTPStatusError as e:
             # Log the error details from OpenRouter if available
             error_detail = f"OpenRouter API Error: {e.response.status_code}"
             try:
                 error_body = e.response.json()
                 error_detail += f" - {error_body.get('error', {}).get('message', e.response.text)}"
             except Exception: # If response is not JSON or parsing fails
                 error_detail += f" - {e.response.text}"
             print(f"HTTPStatusError calling OpenRouter: {error_detail}") # Log the error
             raise HTTPException(status_code=e.response.status_code, detail=error_detail)
        except httpx.RequestError as e:
             print(f"RequestError calling OpenRouter: {e}") # Log the error
             raise HTTPException(status_code=503, detail=f"Could not connect to OpenRouter: {e}")
        except Exception as e: # Catch unexpected errors
             print(f"Unexpected error calling OpenRouter: {e}") # Log the error
             raise HTTPException(status_code=500, detail=f"An unexpected error occurred: {e}")


@router.post("/text", response_model=APIResponse) # Add response_model
async def open_router_text(
    token: Annotated[str, Depends(verify_token)],
    model: ModelName = Query(..., description="Select a model"),
    prompt: str = Query(..., description="Enter your prompt")
):
    payload = {
        "model": model.value, # Use the enum value
        "messages": [{"role": "user", "content": prompt}],
    }
    try:
        response_data = await _call_openrouter(payload)
        if response_data.get("choices") and response_data["choices"][0].get("message"):
            content = response_data["choices"][0]["message"].get("content")
            return APIResponse(success=True, data={"response": content})
        else:
            # Log the unexpected response structure for debugging
            print(f"Unexpected OpenRouter response structure: {response_data}")
            return APIResponse(success=False, error="Unexpected response structure from OpenRouter.", data=response_data)
    except HTTPException as e:
        # Re-raise HTTPException to let FastAPI handle it
        raise e
    except Exception as e:
        # Catch other potential errors during the call or processing
        print(f"Error processing OpenRouter text response: {e}")
        raise HTTPException(status_code=500, detail=f"Error processing response: {str(e)}")


@router.post("/chat-with-memory", response_model=APIResponse)
async def open_router_chat_with_memory(
    request: ChatMemoryRequest,
    token: Annotated[str, Depends(verify_token)]
):
    messages_dict = [msg.dict() for msg in request.messages]

    payload = {
        "model": request.model,
        "messages": messages_dict,
    }
    try:
        response_data = await _call_openrouter(payload)
        if response_data.get("choices") and response_data["choices"][0].get("message"):
            content = response_data["choices"][0]["message"].get("content")
            return APIResponse(success=True, data={"response": content})
        else:
            # Log the unexpected response structure for debugging
            print(f"Unexpected OpenRouter response structure: {response_data}")
            return APIResponse(success=False, error="Unexpected response structure from OpenRouter.", data=response_data)
    except HTTPException as e:
        # Re-raise HTTPException to let FastAPI handle it
        raise e
    except Exception as e:
        # Catch other potential errors during the call or processing
        print(f"Error processing OpenRouter chat-with-memory response: {e}")
        raise HTTPException(status_code=500, detail=f"Error processing response: {str(e)}")


@router.post("/multimodal-url", response_model=APIResponse) # Add response_model
async def open_router_multimodal_url(
    token: Annotated[str, Depends(verify_token)],
    model: MultiModelName = Query(..., description="Select a model"),
    prompt: str = Query(..., description="Enter your prompt (ex: What is in this image?)"),
    image_url: str = Query(..., description="Enter the image URL"),
):
    payload = {
        "model": model.value, # Use the enum value
        "messages": [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": prompt},
                    {"type": "image_url", "image_url": {"url": image_url}},
                ],
            }
        ],
         "max_tokens": 1024, # Example: Set max tokens if needed
    }
    try:
        response_data = await _call_openrouter(payload)
        if response_data.get("choices") and response_data["choices"][0].get("message"):
            content = response_data["choices"][0]["message"].get("content")
            return APIResponse(success=True, data={"response": content})
        else:
             # Log the unexpected response structure for debugging
             print(f"Unexpected OpenRouter response structure: {response_data}")
             return APIResponse(success=False, error="Unexpected response structure from OpenRouter.", data=response_data)
    except HTTPException as e:
        # Re-raise HTTPException to let FastAPI handle it
        raise e
    except Exception as e:
        # Catch other potential errors during the call or processing
        print(f"Error processing OpenRouter multimodal-url response: {e}")
        raise HTTPException(status_code=500, detail=f"Error processing response: {str(e)}")


@router.post("/multimodal-b64", response_model=APIResponse) # Add response_model
async def open_router_multimodal_b64(
    token: Annotated[str, Depends(verify_token)],
    image: UploadFile = File(...),
    model: MultiModelName = Query(..., description="Select a model"),
    prompt: str = Query(..., description="Enter your prompt (ex: What is in this image?)")
):
    try:
        image_bytes = await image.read()
        # Basic validation for file size (e.g., max 10MB)
        if len(image_bytes) > 10 * 1024 * 1024:
             raise HTTPException(status_code=413, detail="Image file size exceeds 10MB limit.")

        encoded_string = base64.b64encode(image_bytes).decode('utf-8')
        # Ensure content_type is provided and valid
        content_type = image.content_type
        if not content_type or not content_type.startswith("image/"):
             raise HTTPException(status_code=400, detail="Invalid or missing image content type.")

        img_data_uri = f"data:{content_type};base64,{encoded_string}"
    except Exception as e:
         raise HTTPException(status_code=400, detail=f"Error processing uploaded image: {e}")


    payload = {
        "model": model.value, # Use the enum value
        "messages": [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": prompt},
                    {"type": "image_url", "image_url": {"url": img_data_uri}},
                ],
            }
        ],
         "max_tokens": 1024, # Example: Set max tokens if needed
    }
    try:
        response_data = await _call_openrouter(payload)
        if response_data.get("choices") and response_data["choices"][0].get("message"):
            content = response_data["choices"][0]["message"].get("content")
            return APIResponse(success=True, data={"response": content})
        else:
             # Log the unexpected response structure for debugging
             print(f"Unexpected OpenRouter response structure: {response_data}")
             return APIResponse(success=False, error="Unexpected response structure from OpenRouter.", data=response_data)
    except HTTPException as e:
        # Re-raise HTTPException to let FastAPI handle it
        raise e
    except Exception as e:
        # Catch other potential errors during the call or processing
        print(f"Error processing OpenRouter multimodal-b64 response: {e}")
        raise HTTPException(status_code=500, detail=f"Error processing response: {str(e)}")


# --- New Endpoint for Rate Limit Status ---
@router.get("/rate-limit-status", response_model=APIResponse)
async def get_rate_limit_status(token: Annotated[str, Depends(verify_token)]):
    """Returns the current OpenRouter API call count for the day."""
    global call_count, last_reset_date, MAX_CALLS_PER_DAY
    today = date.today()

    # Check if reset is needed (although check_and_increment does this too,
    # it's good practice to ensure status reflects the current day)
    if today != last_reset_date:
        call_count = 0
        last_reset_date = today

    return APIResponse(
        success=True,
        data={
            "calls_today": call_count,
            "max_calls_per_day": MAX_CALLS_PER_DAY,
            "limit_resets_on": (last_reset_date + timedelta(days=1)).isoformat() # Show next reset date
        }
    )