from pydantic import BaseModel
from fastapi import FastAPI, Depends, HTTPException
from authx import AuthX, TokenPayload, AuthXConfig

auth_config = AuthXConfig()

auth_config.JWT_ALGORITHM = 'HS256'
auth_config.JWT_SECRET_KEY = 'SECRET_KEY'

auth_config.JWT_TOKEN_LOCATION = ['headers']

app = FastAPI()
security = AuthX(auth_config)


class LoginForm(BaseModel):
    username: str
    password: str


@app.post('/login')
def login(data: LoginForm):
    if data.username == "test" and data.password == "test":
        access_token = security.create_access_token(data.username)
        refresh_token = security.create_refresh_token(data.username)
        return {
            "access_token": access_token,
            "refresh_token": refresh_token
        }
    raise HTTPException(401, "Bad username/password")


@app.post('/refresh')
def refresh(
    refresh_payload: TokenPayload = Depends(security.refresh_token_required)
):
    """
    TODO refresh_token_required error
    """
    access_token = security.create_access_token(refresh_payload.sub)
    return {"access_token": access_token}


@app.get('/protected', dependencies=[Depends(security.access_token_required)])
def protected():
    return "You have access to this protected resource"


if __name__ == '__main__':
    import uvicorn
    uvicorn.run(app, port=8000)