import json
from base64 import b64encode
from functools import partial
from typing import Annotated
from uuid import uuid7

from alembic.testing import requirements
from fastapi import FastAPI, Request, Form, Query
from pydantic import BaseModel, ValidationError
from sqlmodel import select
from starlette.authentication import requires
from starlette.middleware.authentication import AuthenticationMiddleware
from starlette.middleware.sessions import SessionMiddleware
from starlette.responses import HTMLResponse, RedirectResponse
from starlette.templating import Jinja2Templates

import conf
from authentication.authlib import AuthorizationServer, AuthorizationCodeGrant, prepare_oauth_request
from authentication.backend import SessionAuthBackend
from authentication.utils import bcrypt_sha256_verify, get_url_path_with_query
from db.models import User, Session, ClientApplication
from db.session import get_session_context

authentication = FastAPI()
authentication.add_middleware(SessionMiddleware, secret_key=conf.SECRET_KEY)

oauth = FastAPI()
oauth.add_middleware(AuthenticationMiddleware, backend=SessionAuthBackend())
oauth.add_middleware(SessionMiddleware, secret_key=conf.SECRET_KEY)

auth_server = AuthorizationServer()
code_grant = partial(AuthorizationCodeGrant, server=auth_server)

templates = Jinja2Templates(directory='authentication/templates')

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

async def base_login(request: Request, form: LoginForm | None = None, errors: dict[str, str] | None = None):
    return templates.TemplateResponse(request, 'login.html', context={
        'form': form,
        'errors': errors,
    })

@authentication.get('/login/', response_class=HTMLResponse)
async def get_login(request: Request, next: Annotated[str, Query()] = None):
    return await base_login(request)

@authentication.post('/login/', response_class=HTMLResponse)
async def post_login(request: Request, username: Annotated[str, Form()] = '', password: Annotated[str, Form()] = '', next: Annotated[str, Query()] = None):
    form = None
    errors = {}

    try:
        form = LoginForm.model_validate({
            'username': username,
            'password': password
        })
        with get_session_context() as db_session:
            user = db_session.query(User).filter(User.username == form.username).first()

        if not bcrypt_sha256_verify(password, user.password if user is not None else '') or user is None:
            errors.update({'password': 'Invalid username or password'})

        with get_session_context() as db_session:
            session = Session(id=uuid7(), user=user)
            db_session.add(session)

        request.session.update({'id': str(session.id)})

        if next is not None:
            return RedirectResponse(next)

    except ValidationError as exc:
        errors.update({ e['loc'][0]: e['msg'] for e in exc.errors() })

    return await base_login(request, form, errors)

async def base_authorize(request: Request, user: User | None = None):
    if not conf.IS_LOCAL:
        with get_session_context() as db_session:
            db_session.query(Session).filter(Session.id == request.session.get('id')).delete()

    return auth_server.create_authorization_response(request, request.user, grant=None)

@oauth.get('/authorize/')
@requires(['authenticated', 'session'], redirect='get_login')
async def get_authorize(request: Request):
    return await base_authorize(request, request.user)

@oauth.post('/authorize/')
@requires(['authenticated', 'session'], redirect='get_login')
async def post_authorize(request: Request):
    return await base_authorize(request, request.user)

@oauth.post('/token/')
async def post_token(request: Request):
    await prepare_oauth_request(request)
    return auth_server.create_token_response(request)

@authentication.get('/connect/')
async def connect(request: Request):
    origin = request.url.scheme + '://' + request.url.netloc

    with get_session_context() as session:
        app = session.exec(select(ClientApplication)).first()

    return RedirectResponse(
        '/oauth' +  get_url_path_with_query(oauth, 'get_authorize', {
            'client_id': app.client_id,
            'redirect_uri': f'{origin}/oauth' + get_url_path_with_query(oauth, 'callback'),
            'response_type': 'code',
            'state': b64encode(json.dumps({
                'client_id': app.client_id,
            }).encode()).decode(),
        })
    )

@oauth.get('/config/')
async def get_config(request: Request):
    with get_session_context() as session:
        app = session.exec(select(ClientApplication)).first()

    return {
        'client_id': app.client_id,
    }