import os
import requests
from flask import Flask, request, send_file
from selenium import webdriver
from selenium.common.exceptions import WebDriverException
from PIL import Image
from io import BytesIO

# フォントをダウンロードする関数
def download_noto_sans_cjk():
    url = "https://noto-website-2.storage.googleapis.com/pkgs/NotoSansCJKjp-hinted.zip"
    response = requests.get(url)

    # フォントディレクトリを作成
    font_dir = "/usr/share/fonts/truetype/noto"
    os.makedirs(font_dir, exist_ok=True)

    # ダウンロードしたフォントを展開
    with open("/tmp/NotoSansCJKjp.zip", "wb") as f:
        f.write(response.content)

    os.system("unzip /tmp/NotoSansCJKjp.zip -d " + font_dir)
    os.system("fc-cache -fv")  # フォントキャッシュを更新

# フォントのダウンロードとインストールを実行
download_noto_sans_cjk()

app = Flask(__name__)

def take_screenshot(url, screenw=1080, screenh=720):
    options = webdriver.ChromeOptions()
    options.add_argument('--headless')
    options.add_argument('--no-sandbox')
    options.add_argument('--disable-dev-shm-usage')

    try:
        wd = webdriver.Chrome(options=options)
        wd.set_window_size(screenw, screenh)  # Set window size based on screenw and screenh
        wd.get(url)
        wd.implicitly_wait(10)
        screenshot = wd.get_screenshot_as_png()
    except WebDriverException:
        # Return a blank image as a placeholder if there's an error
        return Image.new('RGB', (1, 1))
    finally:
        if wd:
            wd.quit()

    return Image.open(BytesIO(screenshot))

@app.route('/', methods=['GET'])
def screenshot():
    url = request.args.get('url')
    if not url:
        return "URL parameter is required.", 400

    # ウィンドウサイズのパラメーターを取得
    screenw = request.args.get('screenw', default=1080, type=int)
    screenh = request.args.get('screenh', default=720, type=int)
    
    # 画像リサイズのパラメーターを取得
    width = request.args.get('width', type=int)
    height = request.args.get('height', type=int)

    # スクリーンショットを取得
    image = take_screenshot(url, screenw=screenw, screenh=screenh)

    # リサイズ処理
    if width and height:
        # 両方のパラメーターが提供された場合
        image = image.resize((width, height), Image.ANTIALIAS)
    elif width:
        # width のみが提供された場合、アスペクト比に基づいて height を計算
        aspect_ratio = image.height / image.width
        height = int(width * aspect_ratio)
        image = image.resize((width, height), Image.ANTIALIAS)
    elif height:
        # height のみが提供された場合、アスペクト比に基づいて width を計算
        aspect_ratio = image.width / image.height
        width = int(height * aspect_ratio)
        image = image.resize((width, height), Image.ANTIALIAS)

    # 結果画像をメモリに保存してレスポンスとして返却
    img_io = BytesIO()
    image.save(img_io, 'PNG')
    img_io.seek(0)

    return send_file(img_io, mimetype='image/png')

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