from seleniumwire import webdriver
from selenium.webdriver.chrome.options import Options
from fastapi import FastAPI, Request
import uvicorn
import time
import json
from urllib.parse import unquote

app = FastAPI()

@app.get("/")
def main():
    return {"code": 200,"msg":"Success"}
    
@app.get("/chrome")
def chrome(url:str=None,wait:int=5,header:str=None):
    
    target_url = unquote(url)
    wait_time = unquote(wait)
    header_list = unquote(header)

    if target_url == None:
        return {"code": 500,"msg":"No target URL"}

    if wait_time not in range(0, 31):
        return {"code": 500,"msg":"The waiting time must be between 0 and 30"}

    options = Options()
    options.add_argument('--headless')
    
    driver = webdriver.Chrome(options=options)

    try:
        if header_list == None:
            pass
        else:
            header_array = json.loads(header_list)
            driver.header_overrides = header_array
    except ValueError:
        driver.quit()
        return {"code": 500,"msg":"The header field is not JSON"}
    except Exception as e:
        driver.quit()
        return {"code": 500,"msg":"Unable to overwrite header"}
    
    driver.get(target_url)

    print(driver.page_source)
    
    time.sleep(wait_time)

    # 获取当前URL
    current_url = driver.current_url
    
    # 获取页面源代码
    page_source = driver.page_source
    
    # 获取cookie
    cookies = driver.get_cookies()

    # 是否有跳转过
    is_jump = (target_url != current_url)
    
    data = {
        "url": current_url,
        "page_source": page_source,
        "cookies": cookies,
        "is_jump": is_jump
    }
    
    driver.quit()
    return {"code": 200,"data":data}

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