# !/usr/bin/env python3
# -*- encoding : utf-8 -*-
# @Filename    : views.py
# @Software    : VSCode
# @Datetime    : 2023/03/20 15:56:34
# @Author      : leo liu
# @Version     : 1.0
# @Description :

from typing import Optional, Any
from fastapi import APIRouter, Depends, Header
from sqlalchemy.orm.session import Session

from extensions import logger
from core import security
from utils import response_code, custom_exc
from db.session import get_db
from models import User, ChatHistory

from .schemas import chat_schema
from .ctrl.chat import CTRLChat
from ..auth.crud.user import CRUDUser
from .crud.chat import CRUDChat

router = APIRouter()

def get_current_user(
    db: Session = Depends(get_db), token: Optional[str] = Header(None)
) -> User:
    """
    根据header中token 获取当前用户
    :param db:
    :param token:
    :return:
    """
    if not token:
        raise custom_exc.TokenAuthError(err_desc='headers not found token')

    crud_user = CRUDUser(User)

    token_data = security.check_jwt_token(token)
    user = crud_user.get(db, id=token_data.sub)
    if not user:
        raise custom_exc.UserNotFound(err_desc="user not found")
    return user

@router.post("/chat/ask", summary="提问")
async def ask(
    *,
    db: Session = Depends(get_db),
    security: Optional[bool] = Depends(security.check_jwt_token),
    current_user: User = Depends(get_current_user),
    chat_info:chat_schema.Chat = None
) -> Any:
    if not chat_info:
        return response_code.resp_200(data="", message="no chat info")

    crud_chat = CRUDChat(ChatHistory)
    ctrl_chat = CTRLChat()

    q_times = crud_chat.get_chat_count(db, current_user)
    q_limit_day = current_user.q_limit_day

    if ((q_limit_day > 0 and q_times >= q_limit_day) or q_limit_day < 0):
        return response_code.resp_200(data="您今天的提问次数已达到限额，欢迎您明日再来！", message="success chat with j-med")

    openai_keys = crud_chat.get_openai_keys(db)
    default_max_tokens = crud_chat.get_default_max_tokens(db)

    if not openai_keys or len(openai_keys) <= 0:
        return response_code.resp_200(data="", message="no openai_key defined")

    key = openai_keys[0]
    tokens_limit = current_user.tokens_limit

    if tokens_limit == 0:
        tokens_limit = default_max_tokens

    answer = ctrl_chat.chat(key, chat_info.question, tokens_limit)

    crud_chat.update_chat_count(db, current_user)
    crud_chat.new_chat(db, current_user, chat_info.question, answer)

    return response_code.resp_200(data=answer, message="success chat with j-med")