# coding: utf-8
from secrets import compare_digest

from fastapi import Depends, FastAPI, status, Security, HTTPException
from fastapi.security.api_key import APIKeyQuery, APIKeyCookie, APIKeyHeader
from fastapi.openapi.utils import get_openapi
from prefect.server.api.server import create_orion_api, create_ui_app


def get_token(
    access_token_query: str | None = Security(APIKeyQuery(name="token", auto_error=False)),
    access_token_header: str | None = Security(APIKeyHeader(name="token", auto_error=False)),
    access_token_cookie: str | None = Security(APIKeyCookie(name="token", auto_error=False)),
):
    if access_token_query is not None:
        return access_token_query

    elif access_token_header is not None:
        return access_token_header

    elif access_token_cookie is not None:
        return access_token_cookie

    else:
        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="token is missing")


def verify_token(token=Depends(get_token)):
    if compare_digest(token, "secret") is False:
        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="bad token")


custom_app = FastAPI()
api_app = create_orion_api(dependencies=[Depends(verify_token)])
custom_app.mount("/api", api_app)
custom_app.mount("/", create_ui_app(ephemeral=False))


def openapi():
    partial_schema = get_openapi(title="Prefect API customized", version="0.1.0", routes=api_app.routes)
    new_schema = partial_schema.copy()
    new_schema["paths"] = {}
    for path, value in partial_schema["paths"].items():
        new_schema["paths"][f"/api{path}"] = value

    new_schema["info"]["x-logo"] = {"url": "static/prefect-logo-mark-gradient.png"}
    return new_schema


custom_app.openapi = openapi