#!/usr/bin/env python3
import asyncio
import base64
import argparse
from playwright.async_api import async_playwright
from fastapi.responses import JSONResponse
import time
import re, json

import uvicorn
from fastapi import FastAPI, Query, HTTPException
from pydantic import BaseModel, HttpUrl
from typing import Optional
# from screenshot import capture_screenshot


async def capture_screenshot(url, width=1280, height=800, save_path=None):
    """
    Capture a screenshot of a webpage and return as base64 encoded string.
    Args:
        url (str): The URL to capture
        width (int): Viewport width
        height (int): Viewport height

    Returns:
        str: Base64 encoded screenshot data
    """
    timestamp = time.time()
    timestamp = str(timestamp)
    timestamp = re.sub(r"\.", "_", timestamp)

    async with async_playwright() as p:
        browser = await p.chromium.launch()
        page = await browser.new_page(viewport={'width': width, 'height': height})

        try:
            await page.goto(url, wait_until='networkidle')
        except Exception as e:
            await page.goto(url, wait_until='load')

        screenshot_bytes = await page.screenshot(full_page=True, path=save_path)
        await browser.close()

        # Convert to base64
        base64_screenshot = base64.b64encode(screenshot_bytes).decode('utf-8')
        # base64_screenshot = screenshot_bytes

        # 保存图片
        if save_path:
            with open(save_path, 'wb') as f:
                f.write(screenshot_bytes)
            # print(f"Screenshot saved to {args.output}")

        return base64_screenshot

app = FastAPI(title="Screenshot Service")

class Input(BaseModel):
    url: str = ""
    width: int = 1280
    height: int = 800


class ScreenshotResponse(BaseModel):
    url: str
    base64_image: str
    width: int
    height: int



@app.post("/screenshot/")
async def get_screenshot(input: Input):

    try:
        input = json.loads(input.json())
        url = input['url']
        width = input['width']
        height = input['height']
        base64_image = await capture_screenshot(url, width=width, height=height)
        res = {"base64_image":base64_image, "url":url, "width":width, "height":height}


        res = JSONResponse(status_code=200, content=res)

    except Exception as e:
        res = JSONResponse(status_code=500,  content={"message": str(e)} )

    return res

@app.get("/")
async def root():
    return {
        "service": "Screenshot Service",
        "usage": "GET /screenshot?url=https://example.com&width=1280&height=800"
    }

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


    # 对指定网页进行截屏
    # sudo docker run -itd --name playwright -p 19801:19801 -v /home/liuxin/work:/home/work playwright:v1.1  /bin/bash
    # sudo docker exec -it playwright bash
    # cd /home/work/AI_planner/screenshot
    # python web_server.py  # 启动服务