from flask import Flask, render_template, request, send_from_directory
from selenium import webdriver
from flask import Flask, render_template, request, send_file, jsonify
# import os
app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = 'static/uploads/'
import time

def take_screenshot(url):
    # 创建 Chrome 浏览器对象
    options = webdriver.ChromeOptions()
    options.add_argument('--headless')
    options.add_argument('--no-sandbox')
    options.add_argument('--disable-dev-shm-usage')
    driver = webdriver.Chrome(options=options)

    print("开始访问：",url)
    # 访问网页并截图
    driver.maximize_window()
    
    driver.get(url)
    # time.sleep(5)
    driver.implicitly_wait(100)
    file_name = 'static/uploads/screenshot.png'
    driver.save_screenshot(file_name)
    driver.quit()
    # # 裁剪图片
    # image = Image.open(file_name)
    # cropped_image = image.crop((0, 0, 1000, 1000))
    # cropped_image.save(file_name)
    # 返回文件名
    print("返回文件名")
    return file_name


@app.route('/')
def index():
    return render_template('index.html')

@app.route('/screenshot', methods=['POST'])
def screenshot():
    url = request.form['url']
    file_name = take_screenshot(url)
    return jsonify({"previewUrl": file_name, "downloadUrl":file_name})

@app.route('/uploads/<filename>')
def download_file(filename):
    return send_from_directory(app.config['UPLOAD_FOLDER'], filename, as_attachment=True)

if __name__ == '__main__':
    app.run(debug=True)