api.py

import asyncio
import time
import uuid
from typing import List, Optional, Dict, Union, Callable, Iterator, AsyncGenerator
from contextlib import asynccontextmanager
from dataclasses import dataclass
import logging
import os
import io
from google import genai
from google.genai import types
import computer_control_helper
from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, Field
from seleniumbase import Driver
from selenium.webdriver.common.by import By
from selenium.webdriver.common.keys import Keys
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.common.exceptions import TimeoutException, NoSuchElementException, StaleElementReferenceException
import random
from streaming import StreamProcessor, create_streaming_response, StreamConfig, StreamingResponseGenerator

# Virtual display setup for Linux headless environments
import platform
if platform.system() == 'Linux':
    from pyvirtualdisplay import Display
    display = Display(visible=0, size=(1920, 1080))
    display.start()
    logger = logging.getLogger(__name__)
    logger.info("Started virtual display for Linux environment")

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
)
logger = logging.getLogger(__name__)

@dataclass
class Config:
    lmarena_url: str = "https://beta.lmarena.ai/?mode=direct"
    driver_timeout: int = 2
    response_timeout: int = 900
    poll_interval: float = 0.05
    new_chat_click_max_attempts: int = 3
    new_chat_click_retry_delay_seconds: float = 0.1
    new_chat_click_success_pause_seconds: float = 0.5
    page_load_wait_after_refresh_seconds: float = 5
    stabilization_timeout: float = 1.0
    max_inactivity: float = 10.0
config = Config()
logger.info("Configuration loaded.")

class LmArenaError(Exception): pass
class APIError(Exception):
    def __init__(self, message: str, status_code: int):
        self.message = message
        self.status_code = status_code
class ChatInteractionError(APIError):
    def __init__(self, message: str):
        super().__init__(message, status_code=502)
class ModelSelectionError(LmArenaError): pass
class DriverNotAvailableError(LmArenaError): pass

class Message(BaseModel):
    role: str
    content: Union[str, List[Dict[str, str]]]
class ChatCompletionRequest(BaseModel):
    messages: List[Message]
    model: str
    stream: Optional[bool] = False
class Usage(BaseModel):
    prompt_tokens: int; completion_tokens: int; total_tokens: int
class Choice(BaseModel):
    index: int
    message: Optional[Dict[str, str]] = None
    delta: Optional[Dict[str, str]] = None
    finish_reason: Optional[str] = None
class ChatCompletionResponse(BaseModel):
    id: str; object: str; created: int; model: str
    choices: List[Choice]
    usage: Optional[Usage] = None

class ModelInfo(BaseModel):
    id: str
    object: str = "model"
    created: int = Field(default_factory=lambda: int(time.time()))
    owned_by: str = "lmarena"

class ModelListResponse(BaseModel):
    object: str = "list"
    data: List[ModelInfo]


class DriverManager:
    def __init__(self):
        logger.info("DriverManager instance created.")
        self._driver: Optional[Driver] = None
        self._lock = asyncio.Lock()
        self._genai_client = None
        if os.environ.get("GEMINI_API_KEY"):
            try:
                self._genai_client = genai.Client(api_key=os.environ.get("GEMINI_API_KEY"))
                logger.info("Gemini client initialized successfully.")
            except Exception as e:
                logger.error(f"Failed to initialize Gemini client: {e}", exc_info=True)
                self._genai_client = None # Ensure it's None if init fails
        else:
            logger.info("GEMINI_API_KEY not set, Gemini client will not be used for captcha.")

    async def initialize(self) -> None:
        async with self._lock:
            if self._driver is not None:
                logger.warning("Driver initialization called but driver already exists.")
                return

            loop = asyncio.get_event_loop()
            logger.info("Initializing Selenium driver...")

            def _sync_initialize_driver_logic():
                logger.info("Executing synchronous driver initialization and enhanced readiness check.")
                temp_driver = None
                try:
                    temp_driver = Driver(uc=True, headless=False)
                    logger.info("Driver instantiated. Opening URL...")
                    temp_driver.open(config.lmarena_url)
                    logger.info(f"URL '{config.lmarena_url}' opened.")

                    # --- STAGE 1 CAPTCHA HANDLING (Cloudflare / Pre-site) ---
                    logger.info("Attempting to solve initial (Cloudflare-style) captcha with uc_gui_click_captcha()...")
                    temp_driver.uc_gui_click_captcha()
                    logger.info("uc_gui_click_captcha() completed. Main site should be loading now.")
                    # --- END STAGE 1 ---

                    temp_driver.maximize_window()
                    logger.info("Window maximized.")

                    # ---- STAGE 2 CAPTCHA HANDLING (On-site "Verify Human" popup) ----
                    self._perform_sync_captcha_checks(temp_driver)
                    # ---- END STAGE 2 ----
                    return temp_driver
                except Exception as e:
                    logger.error(f"Synchronous driver initialization failed: {e}", exc_info=True)
                    if temp_driver: temp_driver.quit()
                    raise LmArenaError(f"Failed to initialize driver: {e}") from e

            try:
                self._driver = await loop.run_in_executor(None, _sync_initialize_driver_logic)
                logger.info("Selenium driver initialization process completed successfully.")
            except Exception as e:
                logger.error(f"Asynchronous driver initialization failed: {e}", exc_info=True)
                if self._driver: # Check if driver was partially assigned
                    try: 
                        # Ensure self._driver is used if it was assigned before error
                        driver_to_quit = self._driver
                        self._driver = None # Clear it before attempting quit
                        await loop.run_in_executor(None, driver_to_quit.quit)
                    except Exception as quit_e: 
                        logger.error(f"Failed to quit driver after initialization error: {quit_e}")
                else: # If self._driver was never assigned (e.g. error in Driver() call itself)
                    logger.debug("No driver instance to quit after initialization error.")
                
                if isinstance(e, LmArenaError):
                    raise
                raise LmArenaError(f"Failed to initialize driver: {e}") from e

    def _human_like_reload(self, driver: Driver):
        """Human-like page reload with F5/FN+F5 variation"""
        logger.info("Performing human-like page reload")
        
        # Randomly choose between F5 and FN+F5
        if random.random() > 0.5:
            logger.info("Using F5 key")
            body = driver.find_element(By.TAG_NAME, 'body')
            body.send_keys(Keys.F5)
        else:
            logger.info("Using FN+F5 key combination")
            body = driver.find_element(By.TAG_NAME, 'body')
            body.send_keys(Keys.F5)  # Simulate FN+F5 as F5 since FN is hardware specific
        
        # Add random delay to simulate human behavior
        sleep_time = random.uniform(0.5, 2.0)
        time.sleep(sleep_time)
        logger.info(f"Page reloaded after {sleep_time:.2f}s delay")
    
    def _random_mouse_movement(self, driver: Driver):
        logger.info("Performing natural random mouse movement with pyautogui")
        try:
            import pyautogui
            import random
            import math
            import time
            
            # Get screen dimensions
            screen_width, screen_height = pyautogui.size()
            center_x = screen_width // 2
            center_y = screen_height // 2
            
            # Generate random movement pattern (combination of arcs and lines)
            patterns = [
                # Small circles around center
                lambda: [(int(center_x + 50 * math.cos(angle)), int(center_y + 50 * math.sin(angle))) for angle in [2 * math.pi * i / 8 for i in range(8)]],
                
                # Diagonal sweeps
                lambda: [(100, 100), (screen_width-100, screen_height-100), (100, screen_height-100), (screen_width-100, 100)],
                
                # Random walk
                lambda: [(random.randint(100, screen_width-100), random.randint(100, screen_height-100)) for _ in range(5)]
            ]
            
            # Select and execute random pattern
            pattern = random.choice(patterns)()
            total_points = len(pattern)
            duration = 2.0 / total_points  # Total duration ~2 seconds
            
            for i, (x, y) in enumerate(pattern):
                # Add slight randomness to movement speed
                point_duration = duration * random.uniform(0.8, 1.2)
                pyautogui.moveTo(x, y, duration=point_duration)
                
                # Random micro-pauses between movements
                if i < total_points - 1:
                    time.sleep(random.uniform(0.05, 0.15))
            
            logger.info("Natural mouse movement performed successfully")
        except Exception as e:
            logger.warning(f"Random mouse movement failed: {e}", exc_info=True)

    def _perform_sync_captcha_checks(self, driver: Driver):
        # This method is now exclusively for the on-site "Verify Human" popup.
        logger.info("Checking for on-site ('Verify Human') captcha...")
        
        # Add random mouse movements at start
        self._random_mouse_movement(driver)
        
        # First check if textarea is already interactable AND no captcha popup is visible
        try:
            textarea_locator = (By.TAG_NAME, "textarea")
            textarea = WebDriverWait(driver, 5).until(
                EC.element_to_be_clickable(textarea_locator)
            )
            
            # Detect captcha presence with multiple strategies
            captcha_present = False
            # Allow time for captcha to render
            time.sleep(2)
            # 1. Check for common captcha iframes
            try:
                iframes = driver.find_elements(By.TAG_NAME, 'iframe')
                for iframe in iframes:
                    src = iframe.get_attribute('src') or ''
                    if any(keyword in src for keyword in ['api2/anchor', 'api2/bframe', 'recaptcha', 'hcaptcha.com']):
                        if iframe.is_displayed():
                            captcha_present = True
                            logger.info(f"On-site captcha iframe detected with src containing keyword.")
                            break
            except Exception as e:
                logger.debug(f"Error scanning iframes for on-site captcha: {e}")
            # 2. Check for captcha container divs
            if not captcha_present:
                try:
                    containers = driver.find_elements(By.CSS_SELECTOR, '.g-recaptcha, .grecaptcha-badge, .h-captcha')
                    for c in containers:
                        if c.is_displayed():
                            captcha_present = True
                            logger.info("On-site captcha container detected via CSS selector.")
                            break
                except Exception as e:
                    logger.debug(f"Error scanning on-site captcha containers: {e}")
            # 3. Text cues
            if not captcha_present:
                for cue in ['Verify you are human', 'I am not a robot']:
                    try:
                        elem = driver.find_element(By.XPATH, f"//*[contains(text(), '{cue}')]")
                        if elem.is_displayed():
                            captcha_present = True
                            logger.info(f"On-site captcha text cue '{cue}' detected.")
                            break
                    except Exception:
                        pass

            if textarea.is_enabled() and textarea.is_displayed() and not captcha_present:
                logger.info("No on-site captcha detected. Main UI is ready.")
                return
            else:
                logger.info("Textarea not ready or an on-site captcha indicator was found. Proceeding with AI solver.")
                
        except (TimeoutException, NoSuchElementException):
            logger.info("Chat input textarea not interactable. Proceeding with AI captcha solver.")
        except Exception as e:
            logger.warning(f"Unexpected error checking UI state for on-site captcha: {e}", exc_info=True)
        
        if not self._genai_client:
            logger.error("On-site captcha detected, but Gemini client not available. Cannot proceed.")
            raise LmArenaError("AI Captcha solver is required but not configured.")

        # --- AI-based solver for the on-site captcha ---
        try:
            logger.info("Starting visual AI check for on-site captcha.")
            screenshot = computer_control_helper.capture_screen()
            if not screenshot:
                logger.error("Failed to capture screen for AI captcha check.")
                return
            
            img_byte_arr = io.BytesIO()
            screenshot.save(img_byte_arr, format='PNG')
            img_bytes = img_byte_arr.getvalue()

            model_name = "gemini-2.0-flash" 
            logger.info(f"Using Gemini model: {model_name} for on-site captcha detection.")

            contents = [
                types.Content(
                    role="user",
                    parts=[
                        types.Part.from_bytes(mime_type="image/png", data=img_bytes),
                        types.Part.from_text(text="""find the text "Verify you are human". do not give me the coordinates of the text itself - give me the coordinates of the small box to the LEFT of the text. Example response: 
``json
[
  {"box_2d": [504, 151, 541, 170], "label": "box"}
]
``
If you cannot find the checkbox, respond with "No checkbox found".
"""),
                    ]
                ),
            ]
            
            generate_content_config = types.GenerateContentConfig(response_mime_type="text/plain")
            logger.info("Sending screenshot to Gemini API for analysis.")
            response_stream = self._genai_client.models.generate_content_stream(
                model=model_name, 
                contents=contents,
                config=generate_content_config,
            )
            full_response_text = "".join(chunk.text for chunk in response_stream)
            logger.info(f"Received Gemini response for on-site captcha check: {full_response_text}")
            
            if "No checkbox found" in full_response_text:
                logger.info("Gemini indicated no checkbox found for on-site captcha.")
            else:
                parsed_data = computer_control_helper.parse_json_safely(full_response_text)
                click_target = None
                if isinstance(parsed_data, list) and parsed_data:
                    if isinstance(parsed_data[0], dict) and "box_2d" in parsed_data[0]:
                        click_target = parsed_data[0]
                elif isinstance(parsed_data, dict) and "box_2d" in parsed_data:
                    click_target = parsed_data
                
                if click_target:
                    logger.info(f"On-site captcha checkbox found via Gemini. Clicking coordinates: {click_target}")
                    computer_control_helper.perform_click(click_target)
                    time.sleep(3) # Wait after click
                    logger.info("Click performed. Now reloading page as requested for post-AI solve.")
                    self._human_like_reload(driver)
                    time.sleep(config.page_load_wait_after_refresh_seconds)
                    logger.info("Page reloaded after AI captcha solve.")
                else:
                    logger.info("No valid 'box_2d' data found in Gemini response. Reloading as fallback.")
                    self._human_like_reload(driver)
        except Exception as e:
            logger.error(f"An error occurred during AI visual captcha check: {e}", exc_info=True)

    async def cleanup(self) -> None:
        async with self._lock:
            if self._driver:
                logger.info("Cleaning up and quitting Selenium driver...")
                loop = asyncio.get_event_loop()
                driver_to_quit = self._driver
                self._driver = None
                try:
                    await loop.run_in_executor(None, driver_to_quit.quit)
                    logger.info("Driver quit successfully.")
                except Exception as e:
                    logger.error(f"Error during driver cleanup: {e}", exc_info=True)
            else:
                logger.info("Cleanup called but no driver was active.")

    def get_driver(self) -> Driver:
        if self._driver is None:
            logger.error("Attempted to get driver, but it is not available.")
            raise DriverNotAvailableError("Driver not available")
        logger.debug("Driver instance requested and provided.")
        return self._driver
    
    async def _select_model(self, model_id: str) -> None:
        """Select a model on the LmArena page"""
        driver = self.get_driver()
        logger.info(f"Selecting model: {model_id}")

        def _sync_select_model_logic(drv: Driver, m_id: str):
            logger.debug("Starting model selection logic")
            try:
                # Click model dropdown
                logger.debug("Locating model dropdown")
                dropdown_locator = (By.XPATH, "//button[@data-sentry-source-file='select-model.tsx' and @role='combobox']")
                WebDriverWait(drv, config.driver_timeout).until(
                    EC.element_to_be_clickable(dropdown_locator)
                ).click()
                logger.debug("Clicked model dropdown")

                # Enter model name
                logger.debug("Locating model search input")
                search_locator = (By.XPATH, "//input[@placeholder='Search models' and @cmdk-input]")
                search_element = WebDriverWait(drv, config.driver_timeout).until(
                    EC.visibility_of_element_located(search_locator)
                )
                logger.debug("Clearing search input")
                search_element.clear()
                logger.debug(f"Typing model name: {m_id}")
                search_element.send_keys(m_id)
                logger.debug("Sending ENTER key")
                search_element.send_keys(Keys.ENTER)
                logger.info(f"Selected model: {m_id}")

            except (NoSuchElementException, TimeoutException) as e_se_to:
                logger.warning(f"Model selection for '{m_id}' failed due to {type(e_se_to).__name__}.")
                raise ModelSelectionError(f"Failed to select model {m_id}. Original error: {type(e_se_to).__name__}") from e_se_to
            except Exception as e_sync:
                logger.error(f"Model selection failed for {m_id} with an unexpected error: {e_sync}", exc_info=True)
                raise ModelSelectionError(f"Failed to select model {m_id}") from e_sync

        loop = asyncio.get_event_loop()
        try:
            await loop.run_in_executor(None, _sync_select_model_logic, driver, model_id)
        except ModelSelectionError:
            raise
        except Exception as e_exec:
            logger.error(f"Error executing _select_model in executor: {e_exec}", exc_info=True)
            raise ModelSelectionError(f"Failed to select model {model_id} due to executor error: {e_exec}") from e_exec

    def generate_reload_button_location(self, driver: Driver) -> str:
        # This function is not used by the refined captcha logic, but keeping it as it might be used elsewhere.
        logger.info("Generating reload button location with Gemini")
        try:
            # Capture screenshot
            screenshot = computer_control_helper.capture_screen()
            if not screenshot:
                logger.error("Failed to capture screen for reload button detection")
                return "[]"
            
            img_byte_arr = io.BytesIO()
            screenshot.save(img_byte_arr, format='PNG')
            img_bytes = img_byte_arr.getvalue()
            
            model_name = "gemini-2.0-flash"
            contents = [
                types.Content(
                    role="user",
                    parts=[
                        types.Part.from_bytes(mime_type="image/png", data=img_bytes),
                        types.Part.from_text(text="""Find the reload button on the page. It might be labeled with words like "Reload", "Refresh", or have a circular arrow icon. Return the coordinates of the button in the following format:
``json
[
  {"box_2d": [x1, y1, x2, y2], "label": "reload button"}
]
``
If you cannot find the reload button, respond with "No reload button found".
""")
                    ]
                )
            ]
            
            generate_content_config = types.GenerateContentConfig(response_mime_type="text/plain")
            response_stream = self._genai_client.models.generate_content_stream(
                model=model_name, 
                contents=contents,
                config=generate_content_config,
            )
            full_response_text = "".join(chunk.text for chunk in response_stream)
            logger.info(f"Gemini response for reload button: {full_response_text}")
            return full_response_text
        except Exception as e:
            logger.error(f"Error generating reload button location: {e}", exc_info=True)
            return "[]"

driver_manager = DriverManager()

class ChatHandler:
    @staticmethod
    async def send_message_and_stream_response(prompt: str, model_id: str):
        driver = driver_manager.get_driver()
        request_id = str(uuid.uuid4())
        logger.info(f"[{request_id}] Starting chat interaction. Model: '{model_id}'.")
        try:
            if model_id:
                logger.info(f"[{request_id}] Model specified, selecting '{model_id}'.")
                await driver_manager._select_model(model_id)
            sanitized_prompt = ChatHandler._sanitize_for_bmp(prompt)
            logger.info(f"[{request_id}] Sending prompt (first 50 chars): '{sanitized_prompt[:50]}...'")
            await ChatHandler._send_prompt(driver, sanitized_prompt)
            await ChatHandler._handle_agreement_dialog(driver)
            logger.info(f"[{request_id}] Prompt sent. Streaming response...")
            async for chunk in ChatHandler._stream_response(driver, sanitized_prompt, model_id):
                yield chunk
            logger.info(f"[{request_id}] Finished streaming response from browser.")
        except Exception as e:
            logger.error(f"[{request_id}] Chat interaction failed: {e}", exc_info=True)
            raise ChatInteractionError(f"Chat interaction failed: {e}") from e
        finally:
            logger.info(f"[{request_id}] Cleaning up chat session by clicking 'New Chat'.")
            try:
                await ChatHandler._click_new_chat(driver, request_id)
            except Exception as e_cleanup:
                logger.error(f"[{request_id}] Error clicking 'New Chat' during cleanup: {e_cleanup}", exc_info=True)

    @staticmethod
    def _sanitize_for_bmp(text: str) -> str:
        # This function is simple, so logging is omitted unless for debugging
        return ''.join(c for c in text if ord(c) <= 0xFFFF)
    
    @staticmethod
    async def _send_prompt(driver: Driver, prompt: str):
        logger.info("Typing prompt into textarea.")
        loop = asyncio.get_event_loop()
        await loop.run_in_executor(None, lambda: driver.type('textarea', prompt + "\n"))
        logger.info("Prompt submitted.")
    
    @staticmethod
    async def _handle_agreement_dialog(driver: Driver):
        logger.info("Checking for 'Agree' button in dialog.")
        loop = asyncio.get_event_loop()
        clicked = await loop.run_in_executor(None, lambda: driver.click_if_visible("//button[normalize-space()='Agree']"))
        if clicked:
            logger.info("'Agree' button found and clicked.")
        else:
            logger.info("'Agree' button not visible, skipping.")
    
    @staticmethod
    async def _stream_response(driver: Driver, prompt: str, model: str) -> AsyncGenerator[str, None]:
        """Stream response using stabilization-based approach with a corrected locator."""
        try:
            content_container_locator = (By.XPATH, "(//ol[contains(@class, 'flex-col-reverse')]/div[.//h2[starts-with(@id, 'radix-')]])[1]//div[contains(@class, 'grid') and contains(@class, 'pt-4')]")

            WebDriverWait(driver, config.response_timeout).until(
                EC.presence_of_element_located(content_container_locator)
            )
            logger.info("Assistant response container found. Starting to poll for text.")
            
            stream_processor = StreamProcessor(config=StreamConfig(
                poll_interval=config.poll_interval,
                response_timeout=config.response_timeout,
                stabilization_timeout=config.stabilization_timeout,
                max_inactivity=config.max_inactivity
            ))
            
            text_stream = stream_processor.poll_element_text_stream(
                driver=driver,
                element_locator=content_container_locator
            )
            
            stabilized_stream = stream_processor.read_stream_with_stabilization(
                text_stream
            )
            
            async for chunk in ChatHandler._sync_to_async(stabilized_stream):
                yield chunk
                
        except TimeoutException:
            logger.error(f"Streaming error: Timed out waiting for response container to appear.", exc_info=True)
            yield f"\n\nError: Timed out waiting for response from the page."
        except Exception as e:
            logger.error(f"Streaming error: {e}", exc_info=True)
            yield f"\n\nError: {str(e)}"
    
    @staticmethod
    async def _sync_to_async(sync_iter: Iterator[str]) -> AsyncGenerator[str, None]:
        """Convert synchronous iterator to async generator"""
        for item in sync_iter:
            yield item
            await asyncio.sleep(0)
    
    @staticmethod
    async def _click_new_chat(driver: Driver, request_id: str):
        logger.info(f"[{request_id}] Attempting to click 'New Chat' button.")
        loop = asyncio.get_event_loop()
        await loop.run_in_executor(None, lambda: driver.click("//a[contains(@class, 'whitespace-nowrap') and .//h2[contains(text(), 'New Chat')]]"))
        logger.info(f"[{request_id}] 'New Chat' button clicked successfully.")

async def get_available_models() -> List[str]:
    """Scrapes the list of available models from the UI."""
    driver = driver_manager.get_driver()
    loop = asyncio.get_event_loop()

    def _sync_scrape_models(drv: Driver) -> List[str]:
        logger.info("Scraping available models...")
        dropdown_locator = (By.XPATH, "//button[@data-sentry-source-file='select-model.tsx' and @role='combobox']")
        model_item_locator = (By.XPATH, "//div[@cmdk-item and @data-value]")

        try:
            # Click the dropdown to open it
            dropdown_button = WebDriverWait(drv, config.driver_timeout).until(
                EC.element_to_be_clickable(dropdown_locator)
            )
            dropdown_button.click()
            logger.info("Model dropdown clicked.")

            # Wait for the list items to be present
            WebDriverWait(drv, config.driver_timeout).until(
                EC.presence_of_all_elements_located(model_item_locator)
            )
            logger.info("Model list is visible.")
            time.sleep(0.5)  # Brief pause for full render

            # Scrape the model names from the 'data-value' attribute
            model_elements = drv.find_elements(*model_item_locator)
            model_ids = [elem.get_attribute('data-value') for elem in model_elements if elem.get_attribute('data-value')]
            logger.info(f"Found {len(model_ids)} models.")

            # Click the dropdown again to close it
            dropdown_button.click()
            logger.info("Closed model dropdown.")
            
            return model_ids
        except (TimeoutException, NoSuchElementException) as e:
            logger.error(f"Failed to scrape models: {e}", exc_info=True)
            # Attempt to restore state by clicking dropdown again, in case it's stuck open
            try:
                drv.find_element(*dropdown_locator).click()
            except Exception as close_e:
                logger.warning(f"Could not close model dropdown after error: {close_e}")
            raise LmArenaError(f"Could not find or interact with the model dropdown: {e}") from e

    try:
        model_list = await loop.run_in_executor(None, _sync_scrape_models, driver)
        return model_list
    except Exception as e:
        logger.error(f"Error executing model scraping in executor: {e}", exc_info=True)
        raise

@asynccontextmanager
async def lifespan(app: FastAPI):
    logger.info("Application startup sequence initiated.")
    try:
        await driver_manager.initialize()
        logger.info("Application startup sequence completed successfully.")
    except Exception as e:
        logger.critical(f"A critical error occurred during application startup: {e}", exc_info=True)
        # Ensure cleanup is called if initialization fails at any point
        await driver_manager.cleanup() 
        raise
    yield
    logger.info("Application shutdown sequence initiated.")
    await driver_manager.cleanup()
    logger.info("Application shutdown sequence completed.")

app = FastAPI(lifespan=lifespan)

@app.get("/health")
async def health_check():
    logger.info("Health check endpoint called.")
    try:
        driver_manager.get_driver()
        logger.info("Health check status: healthy, driver is available.")
        return {"status": "healthy", "driver": "available"}
    except DriverNotAvailableError:
        logger.warning("Health check status: unhealthy, driver is unavailable.")
        return {"status": "unhealthy", "driver": "unavailable"}

@app.get("/models", response_model=ModelListResponse)
async def list_models():
    """Returns a list of available models."""
    logger.info("Received request for /models endpoint.")
    try:
        model_ids = await get_available_models()
        model_data = [ModelInfo(id=model_id) for model_id in model_ids]
        return ModelListResponse(data=model_data)
    except DriverNotAvailableError:
        logger.error("Models endpoint called but driver is not available.")
        raise HTTPException(status_code=503, detail="Service unavailable: The backend driver is not ready.")
    except Exception as e:
        logger.error(f"An unexpected error occurred while fetching models: {e}", exc_info=True)
        raise HTTPException(status_code=500, detail=f"An unexpected error occurred while fetching models: {str(e)}")

@app.post("/chat/completions", response_model=ChatCompletionResponse)
async def chat_completions(request: ChatCompletionRequest):
    completion_id = f"chatcmpl-{uuid.uuid4().hex}"
    created_timestamp = int(time.time())
    logger.info(f"[{completion_id}] Received chat completion request for model '{request.model}', stream={request.stream}.")
    full_prompt = "\n".join([msg.content for msg in request.messages if isinstance(msg.content, str)])
    try:
        driver_manager.get_driver()
        if request.stream:
            logger.info(f"[{completion_id}] Handling as a streaming request.")
            return StreamingResponse(
                create_streaming_response(
                    completion_id, 
                    created_timestamp, 
                    request.model, 
                    full_prompt,
                    ChatHandler.send_message_and_stream_response
                ),
                media_type="text/event-stream"
            )
        else:
            logger.info(f"[{completion_id}] Handling as a non-streaming request.")
            return await _create_non_streaming_response(
                completion_id, created_timestamp, request.model, full_prompt
            )
    except DriverNotAvailableError as e:
        logger.error(f"[{completion_id}] Service unavailable: The backend driver is not ready. Error: {e}", exc_info=True)
        raise HTTPException(status_code=503, detail="Service unavailable: The backend driver is not ready.")
    except APIError as e:
        logger.error(f"[{completion_id}] API Error occurred: {e.message} (Status: {e.status_code})", exc_info=True)
        raise HTTPException(status_code=e.status_code, detail=e.message)
    except Exception as e:
        logger.error(f"[{completion_id}] An unexpected processing error occurred: {e}", exc_info=True)
        raise HTTPException(status_code=500, detail=f"An unexpected processing error occurred: {str(e)}")

async def _create_non_streaming_response(completion_id: str, created: int, model: str, prompt: str) -> ChatCompletionResponse:
    logger.info(f"[{completion_id}] Creating non-streaming response.")
    try:
        content_parts = [chunk async for chunk in ChatHandler.send_message_and_stream_response(prompt, model)]
        final_content = "".join(content_parts)
        logger.info(f"[{completion_id}] Non-streaming response generated successfully. Content length: {len(final_content)}.")
        return ChatCompletionResponse(
            id=completion_id,
            object="chat.completion",
            created=created,
            model=model,
            choices=[Choice(index=0, message={"role": "assistant", "content": final_content}, finish_reason="stop")],
            usage=Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0)
        )
    except APIError:
        logger.error(f"[{completion_id}] APIError during non-streaming response creation.", exc_info=True)
        raise
    except Exception as e:
        logger.error(f"[{completion_id}] Exception during non-streaming response creation: {e}", exc_info=True)
        raise HTTPException(status_code=500, detail="Error processing non-streaming request.")

if __name__ == "__main__":
    import uvicorn
    logger.info("Starting application...")
    if not os.getenv("GEMINI_API_KEY"):
        logger.error("FATAL: GEMINI_API_KEY environment variable not set. Captcha solving will be disabled.")
        print("ERROR: GEMINI_API_KEY environment variable not set.")
    else:
        logger.info("GEMINI_API_KEY is set.")

    logger.info("Starting Uvicorn server on 0.0.0.0:8000.")
    uvicorn.run(app, host="0.0.0.0", port=8000)