import os
import asyncio
from fastapi import FastAPI, Request, Form
from fastapi.responses import HTMLResponse, JSONResponse, FileResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from helium import start_chrome, kill_browser, go_to
import uuid
import logging
import pathlib
import time

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

# Create FastAPI app
app = FastAPI()

# Determine base directory
BASE_DIR = pathlib.Path(__file__).parent.parent

# Set up templates and static files
templates_path = BASE_DIR / "templates"
static_path = BASE_DIR / "static"
screenshots_path = BASE_DIR / "screenshots"

templates = Jinja2Templates(directory=str(templates_path))
app.mount("/static", StaticFiles(directory=str(static_path)), name="static")

# Create screenshots directory if it doesn't exist
os.makedirs(str(screenshots_path), exist_ok=True)

# Mount the screenshots directory
app.mount("/screenshots", StaticFiles(directory=str(screenshots_path)), name="screenshots")

@app.get("/", response_class=HTMLResponse)
async def read_root(request: Request):
    return templates.TemplateResponse("index.html", {"request": request})

@app.post("/take-screenshot")
async def take_screenshot(url: str = Form(...)):
    logger.info(f"Taking screenshot of URL: {url}")
    
    try:
        # Generate a unique filename
        filename = f"{uuid.uuid4()}.png"
        filepath = str(screenshots_path / filename)
        
        # Log environment info
        logger.info(f"HOME env: {os.environ.get('HOME')}")
        logger.info(f"Current working directory: {os.getcwd()}")
        
        # Launch Chrome with simple options
        logger.info("Launching Chrome browser")
        driver = start_chrome(headless=True)
        
        try:
            # Navigate to URL
            logger.info(f"Navigating to URL: {url}")
            go_to(url)
            
            # Wait for page to load
            logger.info("Waiting for page to load")
            time.sleep(3)
            
            # Take the screenshot
            logger.info(f"Taking screenshot at {filepath}")
            driver.save_screenshot(filepath)
            logger.info(f"Screenshot saved to {filepath}")
        finally:
            # Always close the browser
            kill_browser()
            logger.info("Browser closed")
        
        return JSONResponse({
            "success": True,
            "screenshot_url": f"/screenshots/{filename}"
        })
    
    except Exception as e:
        logger.error(f"Error taking screenshot: {str(e)}")
        return JSONResponse({
            "success": False,
            "error": str(e)
        }, status_code=500)

# Add a health check endpoint
@app.get("/health")
async def health_check():
    return {"status": "ok"}