chat.py

import datetime
from typing import List, Optional, Union

import orjson
import sqlparse
from sqlalchemy import and_, select, update
from sqlalchemy import desc, func
from sqlalchemy.orm import aliased

from apps.chat.models.chat_model import Chat, ChatRecord, CreateChat, ChatInfo, RenameChat, ChatQuestion, ChatLog, \
    TypeEnum, OperationEnum, ChatRecordResult, ChatLogHistory, ChatLogHistoryItem
from apps.datasource.crud.recommended_problem import get_datasource_recommended_chart
from apps.datasource.models.datasource import CoreDatasource
from apps.db.constant import DB
from apps.system.crud.assistant import AssistantOutDs, AssistantOutDsFactory
from apps.system.schemas.system_schema import AssistantOutDsSchema
from common.core.deps import CurrentAssistant, SessionDep, CurrentUser, Trans
from common.utils.utils import extract_nested_json


def get_chat_record_by_id(session: SessionDep, record_id: int):
    record: ChatRecord | None = None

    stmt = select(ChatRecord.id, ChatRecord.question, ChatRecord.chat_id, ChatRecord.datasource, ChatRecord.engine_type,
                  ChatRecord.ai_modal_id, ChatRecord.create_by).where(
        and_(ChatRecord.id == record_id))
    result = session.execute(stmt)
    for r in result:
        record = ChatRecord(id=r.id, question=r.question, chat_id=r.chat_id, datasource=r.datasource,
                            engine_type=r.engine_type, ai_modal_id=r.ai_modal_id, create_by=r.create_by)
    return record


def get_chat(session: SessionDep, chat_id: int) -> Chat:
    statement = select(Chat).where(Chat.id == chat_id)
    chat = session.exec(statement).scalars().first()
    return chat


def list_chats(session: SessionDep, current_user: CurrentUser) -> List[Chat]:
    oid = current_user.oid if current_user.oid is not None else 1
    chart_list = session.query(Chat).filter(and_(Chat.create_by == current_user.id, Chat.oid == oid)).order_by(
        Chat.create_time.desc()).all()
    return chart_list


def list_recent_questions(session: SessionDep, current_user: CurrentUser, datasource_id: int) -> List[str]:
    chat_records = (
        session.query(
            ChatRecord.question
        )
        .join(Chat, ChatRecord.chat_id == Chat.id)  # 关联Chat表
        .filter(
            Chat.datasource == datasource_id,  # 使用Chat表的datasource字段
            ChatRecord.question.isnot(None),
            ChatRecord.create_by == current_user.id
        )
        .group_by(ChatRecord.question)
        .order_by(desc(func.max(ChatRecord.create_time)))
        .limit(10)
        .all()
    )
    return [record[0] for record in chat_records] if chat_records else []


def rename_chat_with_user(session: SessionDep, current_user: CurrentUser, rename_object: RenameChat) -> str:
    chat = session.get(Chat, rename_object.id)
    if not chat:
        raise Exception(f"Chat with id {rename_object.id} not found")
    if chat.create_by != current_user.id:
        raise Exception(f"Chat with id {rename_object.id} not Owned by the current user")
    chat.brief = rename_object.brief.strip()[:20]
    chat.brief_generate = rename_object.brief_generate
    session.add(chat)
    session.flush()
    session.refresh(chat)

    brief = chat.brief
    session.commit()
    return brief


def rename_chat(session: SessionDep, rename_object: RenameChat) -> str:
    chat = session.get(Chat, rename_object.id)
    if not chat:
        raise Exception(f"Chat with id {rename_object.id} not found")

    chat.brief = rename_object.brief.strip()[:20]
    chat.brief_generate = rename_object.brief_generate
    session.add(chat)
    session.flush()
    session.refresh(chat)

    brief = chat.brief
    session.commit()
    return brief


def delete_chat(session, chart_id) -> str:
    chat = session.query(Chat).filter(Chat.id == chart_id).first()
    if not chat:
        return f'Chat with id {chart_id} has been deleted'

    session.delete(chat)
    session.commit()

    return f'Chat with id {chart_id} has been deleted'


def delete_chat_with_user(session, current_user: CurrentUser, chart_id) -> str:
    chat = session.query(Chat).filter(Chat.id == chart_id).first()
    if not chat:
        return f'Chat with id {chart_id} has been deleted'
    if chat.create_by != current_user.id:
        raise Exception(f"Chat with id {chart_id} not Owned by the current user")
    session.delete(chat)
    session.commit()

    return f'Chat with id {chart_id} has been deleted'


def get_chart_config(session: SessionDep, chart_record_id: int):
    stmt = select(ChatRecord.chart).where(and_(ChatRecord.id == chart_record_id))
    res = session.execute(stmt)
    for row in res:
        try:
            return orjson.loads(row.chart)
        except Exception:
            pass
    return {}


def _format_column(column: dict) -> str:
    """格式化单个column字段"""
    value = column.get('value', '')
    name = column.get('name', '')
    if value != name and name:
        return f"{value}({name})"
    return value


def format_chart_fields(chart_info: dict) -> list:
    fields = []

    # 处理 columns
    for column in chart_info.get('columns') or []:
        fields.append(_format_column(column))

    # 处理 axis
    if axis := chart_info.get('axis'):
        # 处理 x 轴
        if x_axis := axis.get('x'):
            fields.append(_format_column(x_axis))

        # 处理 y 轴
        if y_axis := axis.get('y'):
            if isinstance(y_axis, list):
                for column in y_axis:
                    fields.append(_format_column(column))
            else:
                fields.append(_format_column(y_axis))

        # 处理 series
        if series := axis.get('series'):
            fields.append(_format_column(series))

    return [field for field in fields if field]  # 过滤空字符串


def get_last_execute_sql_error(session: SessionDep, chart_id: int):
    stmt = select(ChatRecord.error).where(and_(ChatRecord.chat_id == chart_id)).order_by(
        ChatRecord.create_time.desc()).limit(1)
    res = session.execute(stmt).scalar()
    if res:
        try:
            obj = orjson.loads(res)
            if obj.get('type') and obj.get('type') == 'exec-sql-err':
                return obj.get('traceback')
        except Exception:
            pass

    return None


def format_json_data(origin_data: dict):
    result = {'fields': origin_data.get('fields') if origin_data.get('fields') else []}
    _list = origin_data.get('data') if origin_data.get('data') else []
    data = format_json_list_data(_list)
    result['data'] = data

    return result


def format_json_list_data(origin_data: list[dict]):
    data = []
    for _data in origin_data if origin_data else []:
        _row = {}
        for key, value in _data.items():
            if value is not None:
                # 检查是否为数字且需要特殊处理
                if isinstance(value, (int, float)):
                    # 整数且超过15位 → 转字符串并标记为文本列
                    if isinstance(value, int) and len(str(abs(value))) > 15:
                        value = str(value)
                    # 小数且超过15位有效数字 → 转字符串并标记为文本列
                    elif isinstance(value, float):
                        decimal_str = format(value, '.16f').rstrip('0').rstrip('.')
                        if len(decimal_str) > 15:
                            value = str(value)
            _row[key] = value
        data.append(_row)

    return data


def get_chat_chart_config(session: SessionDep, chat_record_id: int):
    stmt = select(ChatRecord.chart).where(and_(ChatRecord.id == chat_record_id))
    res = session.execute(stmt)
    for row in res:
        try:
            return orjson.loads(row.chart)
        except Exception:
            pass
    return {}


def get_chart_data_with_user(session: SessionDep, current_user: CurrentUser, chat_record_id: int):
    stmt = select(ChatRecord.data).where(and_(ChatRecord.id == chat_record_id, ChatRecord.create_by == current_user.id))
    res = session.execute(stmt)
    for row in res:
        try:
            return orjson.loads(row.data)
        except Exception:
            pass
    return {}


def get_chat_chart_data(session: SessionDep, chat_record_id: int):
    stmt = select(ChatRecord.data).where(and_(ChatRecord.id == chat_record_id))
    res = session.execute(stmt)
    for row in res:
        try:
            return orjson.loads(row.data)
        except Exception:
            pass
    return {}


def get_chat_predict_data_with_user(session: SessionDep, current_user: CurrentUser, chat_record_id: int):
    stmt = select(ChatRecord.predict_data).where(
        and_(ChatRecord.id == chat_record_id, ChatRecord.create_by == current_user.id))
    res = session.execute(stmt)
    for row in res:
        try:
            return orjson.loads(row.predict_data)
        except Exception:
            pass
    return {}


def get_chat_predict_data(session: SessionDep, chat_record_id: int):
    stmt = select(ChatRecord.predict_data).where(and_(ChatRecord.id == chat_record_id))
    res = session.execute(stmt)
    for row in res:
        try:
            return orjson.loads(row.predict_data)
        except Exception:
            pass
    return {}


def get_chat_with_records_with_data(session: SessionDep, chart_id: int, current_user: CurrentUser,
                                    current_assistant: CurrentAssistant) -> ChatInfo:
    return get_chat_with_records(session, chart_id, current_user, current_assistant, True)


dynamic_ds_types = [1, 3]


def get_chat_with_records(session: SessionDep, chart_id: int, current_user: CurrentUser,
                          current_assistant: CurrentAssistant, with_data: bool = False,
                          trans: Trans = None) -> ChatInfo:
    chat = session.get(Chat, chart_id)
    if not chat:
        raise Exception(f"Chat with id {chart_id} not found")
    if chat.create_by != current_user.id:
        raise Exception(f"Chat with id {chart_id} not Owned by the current user")
    chat_info = ChatInfo(**chat.model_dump())

    if current_assistant and current_assistant.type in dynamic_ds_types:
        out_ds_instance = AssistantOutDsFactory.get_instance(current_assistant)
        ds = out_ds_instance.get_ds(chat.datasource, trans)
    else:
        ds = session.get(CoreDatasource, chat.datasource) if chat.datasource else None

    if not ds:
        chat_info.datasource_exists = False
        chat_info.datasource_name = 'Datasource not exist'
    else:
        chat_info.datasource_exists = True
        chat_info.datasource_name = ds.name
        chat_info.ds_type = ds.type

    sql_alias_log = aliased(ChatLog)
    chart_alias_log = aliased(ChatLog)
    analysis_alias_log = aliased(ChatLog)
    predict_alias_log = aliased(ChatLog)

    stmt = (select(ChatRecord.id, ChatRecord.chat_id, ChatRecord.create_time, ChatRecord.finish_time,
                   ChatRecord.question, ChatRecord.sql_answer, ChatRecord.sql,
                   ChatRecord.chart_answer, ChatRecord.chart, ChatRecord.analysis, ChatRecord.predict,
                   ChatRecord.datasource_select_answer, ChatRecord.analysis_record_id, ChatRecord.predict_record_id,
                   ChatRecord.regenerate_record_id,
                   ChatRecord.recommended_question, ChatRecord.first_chat,
                   ChatRecord.finish, ChatRecord.error,
                   sql_alias_log.reasoning_content.label('sql_reasoning_content'),
                   chart_alias_log.reasoning_content.label('chart_reasoning_content'),
                   analysis_alias_log.reasoning_content.label('analysis_reasoning_content'),
                   predict_alias_log.reasoning_content.label('predict_reasoning_content')
                   )
    .outerjoin(sql_alias_log, and_(sql_alias_log.pid == ChatRecord.id,
                                   sql_alias_log.type == TypeEnum.CHAT,
                                   sql_alias_log.operate == OperationEnum.GENERATE_SQL))
    .outerjoin(chart_alias_log, and_(chart_alias_log.pid == ChatRecord.id,
                                     chart_alias_log.type == TypeEnum.CHAT,
                                     chart_alias_log.operate == OperationEnum.GENERATE_CHART))
    .outerjoin(analysis_alias_log, and_(analysis_alias_log.pid == ChatRecord.id,
                                        analysis_alias_log.type == TypeEnum.CHAT,
                                        analysis_alias_log.operate == OperationEnum.ANALYSIS))
    .outerjoin(predict_alias_log, and_(predict_alias_log.pid == ChatRecord.id,
                                       predict_alias_log.type == TypeEnum.CHAT,
                                       predict_alias_log.operate == OperationEnum.PREDICT_DATA))
    .where(and_(ChatRecord.create_by == current_user.id, ChatRecord.chat_id == chart_id)).order_by(
        ChatRecord.create_time))
    if with_data:
        stmt = select(ChatRecord.id, ChatRecord.chat_id, ChatRecord.create_time, ChatRecord.finish_time,
                      ChatRecord.question, ChatRecord.sql_answer, ChatRecord.sql,
                      ChatRecord.chart_answer, ChatRecord.chart, ChatRecord.analysis, ChatRecord.predict,
                      ChatRecord.datasource_select_answer, ChatRecord.analysis_record_id, ChatRecord.predict_record_id,
                      ChatRecord.regenerate_record_id,
                      ChatRecord.recommended_question, ChatRecord.first_chat,
                      ChatRecord.finish, ChatRecord.error, ChatRecord.data, ChatRecord.predict_data).where(
            and_(ChatRecord.create_by == current_user.id, ChatRecord.chat_id == chart_id)).order_by(
            ChatRecord.create_time)

    result = session.execute(stmt).all()
    record_list: list[ChatRecordResult] = []

    # 批量获取所有ChatRecord的token消耗
    record_ids = [row.id for row in result]
    token_usage_map = {}

    if record_ids:
        # 查询所有相关ChatLog的token_usage
        log_stmt = select(ChatLog.pid, ChatLog.token_usage).where(
            and_(
                ChatLog.pid.in_(record_ids),
                ChatLog.local_operation == False,
                ChatLog.token_usage.is_not(None)  # 排除token_usage为空的记录
            )
        )
        log_results = session.execute(log_stmt).all()

        # 按pid分组计算total_tokens总和
        for pid, token_usage in log_results:
            if pid and token_usage is not None:
                tokens_to_add = 0

                if isinstance(token_usage, dict):
                    # 处理字典类型: {"input_tokens": 961, "total_tokens": 1006, "output_tokens": 45}
                    if token_usage:  # 非空字典
                        if "total_tokens" in token_usage:
                            token_value = token_usage["total_tokens"]
                            if isinstance(token_value, (int, float)):
                                tokens_to_add = int(token_value)
                elif isinstance(token_usage, (int, float)):
                    tokens_to_add = int(token_usage)
                if tokens_to_add > 0:
                    if pid not in token_usage_map:
                        token_usage_map[pid] = 0
                    token_usage_map[pid] += tokens_to_add

    for row in result:
        # 计算耗时
        duration = None
        if row.create_time and row.finish_time:
            try:
                time_diff = row.finish_time - row.create_time
                duration = time_diff.total_seconds()  # 转换为秒
            except Exception:
                duration = None

        # 获取token总消耗
        total_tokens = token_usage_map.get(row.id, 0)

        if not with_data:
            record_list.append(
                ChatRecordResult(id=row.id, chat_id=row.chat_id, create_time=row.create_time,
                                 finish_time=row.finish_time,
                                 duration=duration,
                                 total_tokens=total_tokens,
                                 question=row.question, sql_answer=row.sql_answer, sql=row.sql,
                                 chart_answer=row.chart_answer, chart=row.chart,
                                 analysis=row.analysis, predict=row.predict,
                                 datasource_select_answer=row.datasource_select_answer,
                                 analysis_record_id=row.analysis_record_id, predict_record_id=row.predict_record_id,
                                 regenerate_record_id=row.regenerate_record_id,
                                 recommended_question=row.recommended_question, first_chat=row.first_chat,
                                 finish=row.finish, error=row.error,
                                 sql_reasoning_content=row.sql_reasoning_content,
                                 chart_reasoning_content=row.chart_reasoning_content,
                                 analysis_reasoning_content=row.analysis_reasoning_content,
                                 predict_reasoning_content=row.predict_reasoning_content,
                                 ))
        else:
            record_list.append(
                ChatRecordResult(id=row.id, chat_id=row.chat_id, create_time=row.create_time,
                                 finish_time=row.finish_time,
                                 duration=duration,
                                 total_tokens=total_tokens,
                                 question=row.question, sql_answer=row.sql_answer, sql=row.sql,
                                 chart_answer=row.chart_answer, chart=row.chart,
                                 analysis=row.analysis, predict=row.predict,
                                 datasource_select_answer=row.datasource_select_answer,
                                 analysis_record_id=row.analysis_record_id, predict_record_id=row.predict_record_id,
                                 regenerate_record_id=row.regenerate_record_id,
                                 recommended_question=row.recommended_question, first_chat=row.first_chat,
                                 finish=row.finish, error=row.error, data=row.data, predict_data=row.predict_data))

    result = list(map(format_record, record_list))

    for row in result:
        try:
            data_value = row.get('data')
            if data_value is not None:
                row['data'] = format_json_data(data_value)
        except Exception:
            pass

    chat_info.records = result

    return chat_info


def format_record(record: ChatRecordResult):
    _dict = record.model_dump()

    if record.sql_answer and record.sql_answer.strip() != '' and record.sql_answer.strip()[0] == '{' and \
            record.sql_answer.strip()[-1] == '}':
        _obj = orjson.loads(record.sql_answer)
        _dict['sql_answer'] = _obj.get('reasoning_content')
    if record.sql_reasoning_content and record.sql_reasoning_content.strip() != '':
        _dict['sql_answer'] = record.sql_reasoning_content
    if record.chart_answer and record.chart_answer.strip() != '' and record.chart_answer.strip()[0] == '{' and \
            record.chart_answer.strip()[-1] == '}':
        _obj = orjson.loads(record.chart_answer)
        _dict['chart_answer'] = _obj.get('reasoning_content')
    if record.chart_reasoning_content and record.chart_reasoning_content.strip() != '':
        _dict['chart_answer'] = record.chart_reasoning_content
    if record.analysis and record.analysis.strip() != '' and record.analysis.strip()[0] == '{' and \
            record.analysis.strip()[-1] == '}':
        _obj = orjson.loads(record.analysis)
        _dict['analysis_thinking'] = _obj.get('reasoning_content')
        _dict['analysis'] = _obj.get('content')
    if record.analysis_reasoning_content and record.analysis_reasoning_content.strip() != '':
        _dict['analysis_thinking'] = record.analysis_reasoning_content
    if record.predict and record.predict.strip() != '' and record.predict.strip()[0] == '{' and record.predict.strip()[
        -1] == '}':
        _obj = orjson.loads(record.predict)
        _dict['predict'] = _obj.get('reasoning_content')
        _dict['predict_content'] = _obj.get('content')
    if record.predict_reasoning_content and record.predict_reasoning_content.strip() != '':
        _dict['predict'] = record.predict_reasoning_content
    if record.data and record.data.strip() != '':
        try:
            _obj = orjson.loads(record.data)
            _dict['data'] = _obj
        except Exception:
            pass
    if record.predict_data and record.predict_data.strip() != '':
        try:
            _obj = orjson.loads(record.predict_data)
            _dict['predict_data'] = _obj
        except Exception:
            pass
    if record.sql and record.sql.strip() != '':
        try:
            _dict['sql'] = sqlparse.format(record.sql, reindent=True)
        except Exception:
            pass

    # 格式化duration字段，保留2位小数
    if 'duration' in _dict and _dict['duration'] is not None:
        try:
            # 可以格式化为更易读的形式
            _dict['duration'] = round(_dict['duration'], 2)  # 保留2位小数
        except Exception:
            pass

    # 格式化total_tokens字段
    if 'total_tokens' in _dict and _dict['total_tokens'] is not None:
        try:
            # 确保是整数类型
            _dict['total_tokens'] = int(_dict['total_tokens']) if _dict['total_tokens'] else 0
        except Exception:
            _dict['total_tokens'] = 0

    # 去除返回前端多余的字段
    _dict.pop('sql_reasoning_content', None)
    _dict.pop('chart_reasoning_content', None)
    _dict.pop('analysis_reasoning_content', None)
    _dict.pop('predict_reasoning_content', None)

    return _dict


def get_chat_log_history(session: SessionDep, chat_record_id: int, current_user: CurrentUser, without_steps: bool = False) -> ChatLogHistory:
    """
    获取ChatRecord的详细历史记录

    Args:
        session: 数据库会话
        chat_record_id: ChatRecord的ID
        current_user: 当前用户
        without_steps

    Returns:
        ChatLogHistory: 包含历史步骤和时间信息的对象
    """
    # 1. 首先验证ChatRecord存在且属于当前用户
    chat_record = session.get(ChatRecord, chat_record_id)
    if not chat_record:
        raise Exception(f"ChatRecord with id {chat_record_id} not found")

    if chat_record.create_by != current_user.id:
        raise Exception(f"ChatRecord with id {chat_record_id} not owned by the current user")

    # 2. 查询与该ChatRecord相关的所有ChatLog记录
    chat_logs = session.query(ChatLog).filter(
        ChatLog.pid == chat_record_id,
        ChatLog.operate != OperationEnum.GENERATE_RECOMMENDED_QUESTIONS
    ).order_by(ChatLog.start_time).all()

    # 3. 计算总的时间和token信息
    total_tokens = 0
    steps = []

    for log in chat_logs:
        # 计算单条记录的token消耗
        log_tokens = 0
        if log.token_usage is not None:
            if isinstance(log.token_usage, dict):
                if log.token_usage and "total_tokens" in log.token_usage:
                    token_value = log.token_usage["total_tokens"]
                    if isinstance(token_value, (int, float)):
                        log_tokens = int(token_value)
            elif isinstance(log.token_usage, (int, float)):
                log_tokens = log.token_usage

        # 累加到总token消耗
        total_tokens += log_tokens

        if not without_steps:
            # 计算单条记录的耗时
            duration = None
            if log.start_time and log.finish_time:
                try:
                    time_diff = log.finish_time - log.start_time
                    duration = round(time_diff.total_seconds(), 2)
                except Exception:
                    duration = None

            # 获取操作类型的枚举名称
            operate_name = None
            if log.operate:
                # 如果是OperationEnum枚举实例
                if isinstance(log.operate, OperationEnum):
                    operate_name = log.operate.name
                # 如果是字符串，尝试从枚举值获取名称
                elif isinstance(log.operate, str):
                    try:
                        # 通过枚举值找到对应的枚举实例
                        for enum_item in OperationEnum:
                            if enum_item.value == log.operate:
                                operate_name = enum_item.name
                                break
                    except Exception:
                        operate_name = log.operate
                else:
                    operate_name = str(log.operate)

            # 创建ChatLogHistoryItem
            history_item = ChatLogHistoryItem(
                start_time=log.start_time,
                finish_time=log.finish_time,
                duration=duration,
                total_tokens=log_tokens,
                operate=operate_name,
                local_operation=log.local_operation
            )

            steps.append(history_item)

    # 4. 计算总耗时（使用ChatRecord的时间）
    total_duration = None
    if chat_record.create_time and chat_record.finish_time:
        try:
            time_diff = chat_record.finish_time - chat_record.create_time
            total_duration = round(time_diff.total_seconds(), 2)
        except Exception:
            total_duration = None

    # 5. 创建并返回ChatLogHistory对象
    chat_log_history = ChatLogHistory(
        start_time=chat_record.create_time,  # 使用ChatRecord的create_time
        finish_time=chat_record.finish_time,  # 使用ChatRecord的finish_time
        duration=total_duration,
        total_tokens=total_tokens,
        steps=steps
    )

    return chat_log_history

def get_chat_brief_generate(session: SessionDep, chat_id: int):
    chat = get_chat(session=session, chat_id=chat_id)
    if chat is not None and chat.brief_generate is not None:
        return chat.brief_generate
    else:
        return False


def list_generate_sql_logs(session: SessionDep, chart_id: int) -> List[ChatLog]:
    stmt = select(ChatLog).where(
        and_(ChatLog.pid.in_(select(ChatRecord.id).where(and_(ChatRecord.chat_id == chart_id))),
             ChatLog.type == TypeEnum.CHAT, ChatLog.operate == OperationEnum.GENERATE_SQL)).order_by(
        ChatLog.start_time)
    result = session.execute(stmt).all()
    _list = []
    for row in result:
        for r in row:
            _list.append(ChatLog(**r.model_dump()))
    return _list


def list_generate_chart_logs(session: SessionDep, chart_id: int) -> List[ChatLog]:
    stmt = select(ChatLog).where(
        and_(ChatLog.pid.in_(select(ChatRecord.id).where(and_(ChatRecord.chat_id == chart_id))),
             ChatLog.type == TypeEnum.CHAT, ChatLog.operate == OperationEnum.GENERATE_CHART)).order_by(
        ChatLog.start_time)
    result = session.execute(stmt).all()
    _list = []
    for row in result:
        for r in row:
            _list.append(ChatLog(**r.model_dump()))
    return _list


def create_chat(session: SessionDep, current_user: CurrentUser, create_chat_obj: CreateChat,
                require_datasource: bool = True, current_assistant: CurrentAssistant = None) -> ChatInfo:
    if not create_chat_obj.datasource and require_datasource:
        raise Exception("Datasource cannot be None")

    if not create_chat_obj.question or create_chat_obj.question.strip() == '':
        create_chat_obj.question = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")

    chat = Chat(create_time=datetime.datetime.now(),
                create_by=current_user.id,
                oid=current_user.oid if current_user.oid is not None else 1,
                brief=create_chat_obj.question.strip()[:20],
                origin=create_chat_obj.origin if create_chat_obj.origin is not None else 0)
    ds: CoreDatasource | AssistantOutDsSchema | None = None
    if create_chat_obj.datasource:
        chat.datasource = create_chat_obj.datasource
        if current_assistant and current_assistant.type == 1:
            out_ds_instance: AssistantOutDs = AssistantOutDsFactory.get_instance(current_assistant)
            ds = out_ds_instance.get_ds(chat.datasource)
            ds.type_name = DB.get_db(ds.type)
        else:
            ds = session.get(CoreDatasource, create_chat_obj.datasource)
            if ds.oid != current_user.oid:
                raise Exception(f"Datasource with id {create_chat_obj.datasource} does not belong to current workspace")

        if not ds:
            raise Exception(f"Datasource with id {create_chat_obj.datasource} not found")

        chat.engine_type = ds.type_name
    else:
        chat.engine_type = ''

    chat_info = ChatInfo(**chat.model_dump())

    session.add(chat)
    session.flush()
    session.refresh(chat)
    chat_info.id = chat.id
    session.commit()

    if ds:
        chat_info.datasource_exists = True
        chat_info.datasource_name = ds.name
        chat_info.ds_type = ds.type

    if require_datasource and ds:
        # generate first empty record
        record = ChatRecord()
        record.chat_id = chat.id
        record.datasource = ds.id
        record.engine_type = ds.type_name
        record.first_chat = True
        record.finish = True
        record.create_time = datetime.datetime.now()
        record.create_by = current_user.id
        if isinstance(ds, CoreDatasource) and ds.recommended_config == 2:
            questions = get_datasource_recommended_chart(session, ds.id)
            record.recommended_question = orjson.dumps(questions).decode()
            record.recommended_question_answer = orjson.dumps({
                "content": questions
            }).decode()

        _record = ChatRecord(**record.model_dump())

        session.add(record)
        session.flush()
        session.refresh(record)
        _record.id = record.id
        session.commit()

        chat_info.records.append(_record)

    return chat_info


def save_question(session: SessionDep, current_user: CurrentUser, question: ChatQuestion) -> ChatRecord:
    if not question.chat_id:
        raise Exception("ChatId cannot be None")
    if not question.question or question.question.strip() == '':
        raise Exception("Question cannot be Empty")

    # chat = session.query(Chat).filter(Chat.id == question.chat_id).first()
    chat: Chat = session.get(Chat, question.chat_id)
    if not chat:
        raise Exception(f"Chat with id {question.chat_id} not found")

    record = ChatRecord()
    record.question = question.question
    record.chat_id = chat.id
    record.create_time = datetime.datetime.now()
    record.create_by = current_user.id
    record.datasource = chat.datasource
    record.engine_type = chat.engine_type
    record.ai_modal_id = question.ai_modal_id
    record.regenerate_record_id = question.regenerate_record_id

    result = ChatRecord(**record.model_dump())

    session.add(record)
    session.flush()
    session.refresh(record)
    result.id = record.id
    session.commit()

    return result


def save_analysis_predict_record(session: SessionDep, base_record: ChatRecord, action_type: str) -> ChatRecord:
    record = ChatRecord()
    record.question = base_record.question
    record.chat_id = base_record.chat_id
    record.datasource = base_record.datasource
    record.engine_type = base_record.engine_type
    record.ai_modal_id = base_record.ai_modal_id
    record.create_time = datetime.datetime.now()
    record.create_by = base_record.create_by
    record.chart = base_record.chart
    record.data = base_record.data

    if action_type == 'analysis':
        record.analysis_record_id = base_record.id
    elif action_type == 'predict':
        record.predict_record_id = base_record.id

    result = ChatRecord(**record.model_dump())

    session.add(record)
    session.flush()
    session.refresh(record)
    result.id = record.id
    session.commit()

    return result


def start_log(session: SessionDep, ai_modal_id: int = None, ai_modal_name: str = None, operate: OperationEnum = None,
              record_id: int = None, full_message: Union[list[dict], dict] = None,
              local_operation: bool = False) -> ChatLog:
    log = ChatLog(type=TypeEnum.CHAT, operate=operate, pid=record_id, ai_modal_id=ai_modal_id, base_modal=ai_modal_name,
                  messages=full_message, start_time=datetime.datetime.now(), local_operation=local_operation)

    result = ChatLog(**log.model_dump())

    session.add(log)
    session.flush()
    session.refresh(log)
    result.id = log.id
    session.commit()

    return result


def end_log(session: SessionDep, log: ChatLog, full_message: Union[list[dict], dict, str],
            reasoning_content: str = None,
            token_usage=None) -> ChatLog:
    if token_usage is None:
        token_usage = {}
    log.messages = full_message
    log.token_usage = token_usage
    log.finish_time = datetime.datetime.now()
    log.reasoning_content = reasoning_content if reasoning_content and len(reasoning_content.strip()) > 0 else None

    stmt = update(ChatLog).where(and_(ChatLog.id == log.id)).values(
        messages=log.messages,
        token_usage=log.token_usage,
        finish_time=log.finish_time,
        reasoning_content=log.reasoning_content
    )
    session.execute(stmt)
    session.commit()

    return log


def save_sql_answer(session: SessionDep, record_id: int, answer: str) -> ChatRecord:
    if not record_id:
        raise Exception("Record id cannot be None")

    stmt = update(ChatRecord).where(and_(ChatRecord.id == record_id)).values(
        sql_answer=answer,
    )

    session.execute(stmt)

    session.commit()

    record = get_chat_record_by_id(session, record_id)

    return record


def save_analysis_answer(session: SessionDep, record_id: int, answer: str = '') -> ChatRecord:
    if not record_id:
        raise Exception("Record id cannot be None")

    stmt = update(ChatRecord).where(and_(ChatRecord.id == record_id)).values(
        analysis=answer,
    )

    session.execute(stmt)

    session.commit()

    record = get_chat_record_by_id(session, record_id)

    return record


def save_predict_answer(session: SessionDep, record_id: int, answer: str) -> ChatRecord:
    if not record_id:
        raise Exception("Record id cannot be None")

    stmt = update(ChatRecord).where(and_(ChatRecord.id == record_id)).values(
        predict=answer,
    )

    session.execute(stmt)

    session.commit()

    record = get_chat_record_by_id(session, record_id)

    return record


def save_select_datasource_answer(session: SessionDep, record_id: int, answer: str,
                                  datasource: int = None, engine_type: str = None) -> ChatRecord:
    if not record_id:
        raise Exception("Record id cannot be None")
    record = get_chat_record_by_id(session, record_id)

    record.datasource_select_answer = answer

    if datasource:
        record.datasource = datasource
        record.engine_type = engine_type

    result = ChatRecord(**record.model_dump())

    if datasource:
        stmt = update(ChatRecord).where(and_(ChatRecord.id == record.id)).values(
            datasource_select_answer=record.datasource_select_answer,
            datasource=record.datasource,
            engine_type=record.engine_type,
        )
    else:
        stmt = update(ChatRecord).where(and_(ChatRecord.id == record.id)).values(
            datasource_select_answer=record.datasource_select_answer,
        )

    session.execute(stmt)

    session.commit()

    return result


def save_recommend_question_answer(session: SessionDep, record_id: int,
                                   answer: dict = None, articles_number: Optional[int] = 4) -> ChatRecord:
    if not record_id:
        raise Exception("Record id cannot be None")

    recommended_question_answer = orjson.dumps(answer).decode()

    json_str = '[]'
    if answer and answer.get('content') and answer.get('content') != '':
        try:
            json_str = extract_nested_json(answer.get('content'))

            if not json_str:
                json_str = '[]'
        except Exception as e:
            pass
    recommended_question = json_str

    stmt = update(ChatRecord).where(and_(ChatRecord.id == record_id)).values(
        recommended_question_answer=recommended_question_answer,
        recommended_question=recommended_question,
    )

    session.execute(stmt)
    session.commit()

    record = get_chat_record_by_id(session, record_id)
    record.recommended_question_answer = recommended_question_answer
    record.recommended_question = recommended_question
    if articles_number > 4:
        stmt_chat = update(Chat).where(and_(Chat.id == record.chat_id)).values(
            recommended_question_answer=recommended_question_answer,
            recommended_question=recommended_question,
            recommended_generate=True
        )
        session.execute(stmt_chat)
        session.commit()

    return record


def save_sql(session: SessionDep, record_id: int, sql: str) -> ChatRecord:
    if not record_id:
        raise Exception("Record id cannot be None")

    record = get_chat_record_by_id(session, record_id)

    record.sql = sql

    result = ChatRecord(**record.model_dump())

    stmt = update(ChatRecord).where(and_(ChatRecord.id == record.id)).values(
        sql=record.sql
    )

    session.execute(stmt)

    session.commit()

    return result


def save_chart_answer(session: SessionDep, record_id: int, answer: str) -> ChatRecord:
    if not record_id:
        raise Exception("Record id cannot be None")

    stmt = update(ChatRecord).where(and_(ChatRecord.id == record_id)).values(
        chart_answer=answer,
    )

    session.execute(stmt)

    session.commit()

    record = get_chat_record_by_id(session, record_id)

    return record


def save_chart(session: SessionDep, record_id: int, chart: str) -> ChatRecord:
    if not record_id:
        raise Exception("Record id cannot be None")
    record = get_chat_record_by_id(session, record_id)

    record.chart = chart

    result = ChatRecord(**record.model_dump())

    stmt = update(ChatRecord).where(and_(ChatRecord.id == record.id)).values(
        chart=record.chart
    )

    session.execute(stmt)

    session.commit()

    return result


def save_predict_data(session: SessionDep, record_id: int, data: str = '') -> ChatRecord:
    if not record_id:
        raise Exception("Record id cannot be None")
    record = get_chat_record_by_id(session, record_id)

    record.predict_data = data

    result = ChatRecord(**record.model_dump())

    stmt = update(ChatRecord).where(and_(ChatRecord.id == record.id)).values(
        predict_data=record.predict_data
    )

    session.execute(stmt)

    session.commit()

    return result


def save_error_message(session: SessionDep, record_id: int, message: str) -> ChatRecord:
    if not record_id:
        raise Exception("Record id cannot be None")
    record = get_chat_record_by_id(session, record_id)

    record.error = message
    record.finish = True
    record.finish_time = datetime.datetime.now()

    result = ChatRecord(**record.model_dump())

    stmt = update(ChatRecord).where(and_(ChatRecord.id == record.id)).values(
        error=record.error,
        finish=record.finish,
        finish_time=record.finish_time
    )

    session.execute(stmt)

    session.commit()

    # log error finish
    stmt = update(ChatLog).where(and_(ChatLog.pid == record.id, ChatLog.finish_time.is_(None))).values(
        finish_time=record.finish_time
    )
    session.execute(stmt)
    session.commit()

    return result


def save_sql_exec_data(session: SessionDep, record_id: int, data: str) -> ChatRecord:
    if not record_id:
        raise Exception("Record id cannot be None")
    record = get_chat_record_by_id(session, record_id)

    record.data = data

    result = ChatRecord(**record.model_dump())

    stmt = update(ChatRecord).where(and_(ChatRecord.id == record.id)).values(
        data=record.data,
    )

    session.execute(stmt)

    session.commit()

    return result


def finish_record(session: SessionDep, record_id: int) -> ChatRecord:
    if not record_id:
        raise Exception("Record id cannot be None")
    record = get_chat_record_by_id(session, record_id)

    record.finish = True
    record.finish_time = datetime.datetime.now()

    result = ChatRecord(**record.model_dump())

    stmt = update(ChatRecord).where(and_(ChatRecord.id == record.id)).values(
        finish=record.finish,
        finish_time=record.finish_time
    )

    session.execute(stmt)

    session.commit()

    return result


def get_old_questions(session: SessionDep, datasource: int):
    records = []
    if not datasource:
        return records
    stmt = select(ChatRecord.question).where(
        and_(ChatRecord.datasource == datasource, ChatRecord.question.isnot(None),
             ChatRecord.error.is_(None))).order_by(
        ChatRecord.create_time.desc()).limit(20)
    result = session.execute(stmt)
    for r in result:
        records.append(r.question)
    return records