from collections.abc import Generator
from typing import Annotated, Optional
from fastapi import Depends, Security, HTTPException, status
from fastapi.security.api_key import APIKeyHeader
from sqlmodel import Session, select

from app.core.db import engine
from app.core.settings import ENABLE_API_KEY
from app.models import ApiKey


def get_db_session() -> Generator[Session, None, None]:
    with Session(engine) as session:
        yield session


SessionDep = Annotated[Session, Depends(get_db_session)]


# TODO: after we force the user to provide an API key in the header, change auto_error to True
api_key_header = APIKeyHeader(name="x-api-key", auto_error=False)


def check_api_key(api_key_header: str = Security(api_key_header)) -> Optional[ApiKey]:
    if not ENABLE_API_KEY:
        return None

    # Do not use `SessionDep`` here for now, because we don't want to create a new session
    # for every request that uses an API key.
    # https://pingcap.sentry.io/issues/5476748822
    with Session(engine) as session:
        statement = select(ApiKey).where(ApiKey.key == api_key_header)
        api_key = session.exec(statement).first()
        if not api_key or not api_key.is_active:
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid API key"
            )
        return api_key


ApiKeyDep = Annotated[ApiKey, Depends(check_api_key)]