import os
import time
import base64
import httpx
import asyncio
import uvicorn
from fastapi import FastAPI
from pydantic import BaseModel
from playwright.async_api import async_playwright
from contextlib import asynccontextmanager

# --- Configuration ---
IMGBB_API_KEY = os.getenv("IMGBB_API_KEY", "5765c99f7a977515ba320a689330e508") 
SAVE_FOLDER = "/tmp/downloaded_images"
os.makedirs(SAVE_FOLDER, exist_ok=True)

class GeminiBot:
    def __init__(self):
        self.playwright = None
        self.context = None
        self.page = None
        self.history = {"last_text": "", "last_img_id": ""}

    async def start(self):
        print("Starting Playwright...")
        self.playwright = await async_playwright().start()
        
        # HF Spaces కోసం Chromium + Headless వాడాలి
        self.context = await self.playwright.chromium.launch_persistent_context(
            user_data_dir="./user_data", # Local folder for session
            channel='chromium',          # msedge కాదు, chromium వాడాలి
            headless=True,               # True అయి ఉండాలి (Server కోసం)
            accept_downloads=True,
            args=[
                "--disable-blink-features=AutomationControlled",
                "--no-sandbox",
                "--disable-setuid-sandbox",
                "--disable-dev-shm-usage"
            ]
        )
        self.page = self.context.pages[0]
        
        print("Navigating to Gemini...")
        await self.page.goto("https://gemini.google.com")
        
        try:
            # లాగిన్ అయ్యాక మాత్రమే ఈ ఎలిమెంట్ కనిపిస్తుంది
            await self.page.wait_for_selector("div[contenteditable='true']", timeout=60000)
            print("--- Gemini Bot is ready (Logged In) ---")
        except Exception as e:
            print(f"Error: Login failed or selector not found. {e}")
            await self.page.screenshot(path="debug_error.png")
            # లాగిన్ ఫెయిల్ అయితే కూడా app run avvali kabatti raise cheyyakunda continue cheyyachu
            # కానీ chat work avvakapovachu.

    async def upload_to_imgbb(self, image_path):
        if not IMGBB_API_KEY:
            return None
        try:
            if not image_path or not os.path.exists(image_path): 
                return None
            with open(image_path, "rb") as file:
                img_base64 = base64.b64encode(file.read()).decode('utf-8')
            
            async with httpx.AsyncClient() as client:
                response = await client.post(
                    "https://api.imgbb.com/1/upload", 
                    data={"key": IMGBB_API_KEY, "image": img_base64}
                )
                if response.status_code == 200:
                    return response.json()['data']['url']
                else:
                    return None
        except Exception as e:
            print(f"--- [ERROR] Upload: {e} ---")
            return None

    async def download_image_expert(self):
        try:
            all_imgs = self.page.locator("img[src^='blob:']")
            count = await all_imgs.count()
            if count > 0:
                target_img = all_imgs.last
                img_src = await target_img.get_attribute("src")
                
                if img_src == self.history["last_img_id"]: 
                    return None
                
                box = await target_img.bounding_box()
                if box:
                    await self.page.mouse.move(box['x'] + (box['width'] / 2), box['y'] + (box['height'] / 2))
                    await asyncio.sleep(2) 
                    
                    download_btn = self.page.locator("button:has(mat-icon[fonticon='download'])").last
                    if await download_btn.is_visible():
                        async with self.page.expect_download(timeout=30000) as download_info:
                            try:
                                await download_btn.click(force=True)
                            except:
                                await self.page.evaluate("el => el.click()", await download_btn.element_handle())
                        
                        downloaded_file = await download_info.value
                        file_name = f"gemini_{int(time.time())}.png"
                        save_path = os.path.join(SAVE_FOLDER, file_name)
                        await downloaded_file.save_as(save_path)
                        
                        self.history["last_img_id"] = img_src
                        return save_path
        except Exception as e:
            print(f"--- [ERROR] Download: {e} ---")
        return None

    async def process_chat(self, message: str, file_paths: list = None):
        try:
            # File Upload Logic (Optional - needs adjustment for server paths)
            if file_paths:
                print("File upload requested but not fully implemented for server paths yet.")

            # 1. Message Send
            chat_box = self.page.locator("div[contenteditable='true']").first
            await chat_box.fill(message)
            await self.page.keyboard.press("Enter")
            
            # 2. Wait for Response
            stop_btn = self.page.locator("mat-icon[fonticon='stop']")
            await self.page.wait_for_selector("message-content", timeout=60000)
            
            while await stop_btn.is_visible():
                await asyncio.sleep(0.5)
            
            await asyncio.sleep(2) 
            
            # 3. Get Text
            current_text = await self.page.evaluate("""
                () => {
                    const msgs = document.querySelectorAll('message-content');
                    if (msgs.length > 0) return msgs[msgs.length - 1].innerText;
                    return '';
                }
            """)
            
            # 4. Handle Image
            img_path = await self.download_image_expert()
            img_url = await self.upload_to_imgbb(img_path) if img_path else None
            
            return {"text": current_text, "image_url": img_url}
        except Exception as e:
            return {"text": f"Error: {str(e)}", "image_url": None}

bot = GeminiBot()

@asynccontextmanager
async def lifespan(app: FastAPI):
    await bot.start()
    yield
    if bot.context:
        await bot.context.close()
    if bot.playwright:
        await bot.playwright.stop()

app = FastAPI(lifespan=lifespan)

class ChatRequest(BaseModel):
    message: str
    files: list = []

@app.get("/")
def health_check():
    return {"status": "online", "message": "Gemini API is running"}

@app.post("/chat")
async def chat_api(request: ChatRequest):
    result = await bot.process_chat(request.message, request.files)
    return {"status": "success", "data": result}

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