# coding=utf-8

import time
from functools import reduce
from typing import List, Dict

from django.db.models import QuerySet
from langchain_core.messages import BaseMessage, HumanMessage, SystemMessage, AIMessage

from application.flow.i_step_node import NodeResult, INode
from application.flow.step_node.video_understand_step_node.i_video_understand_node import IVideoUnderstandNode
from application.flow.tools import Reasoning
from knowledge.models import File
from models_provider.tools import get_model_instance_by_model_workspace_id


def _write_context(node_variable: Dict, workflow_variable: Dict, node: INode, workflow, answer: str,
                   reasoning_content: str):
    chat_model = node_variable.get('chat_model')
    message_tokens = node_variable['usage_metadata']['output_tokens'] if 'usage_metadata' in node_variable else 0
    answer_tokens = chat_model.get_num_tokens(answer)
    node.context['message_tokens'] = message_tokens
    node.context['answer_tokens'] = answer_tokens
    node.context['answer'] = answer
    node.context['history_message'] = node_variable['history_message']
    node.context['question'] = node_variable['question']
    node.context['run_time'] = time.time() - node.context['start_time']
    node.context['reasoning_content'] = reasoning_content
    if workflow.is_result(node, NodeResult(node_variable, workflow_variable)):
        node.answer_text = answer


def write_context_stream(node_variable: Dict, workflow_variable: Dict, node: INode, workflow):
    """
    写入上下文数据 (流式)
    @param node_variable:      节点数据
    @param workflow_variable:  全局数据
    @param node:               节点
    @param workflow:           工作流管理器
    """
    response = node_variable.get('result')
    answer = ''
    reasoning_content = ''
    model_setting = node.context.get('model_setting',
                                     {'reasoning_content_enable': False, 'reasoning_content_end': '</think>',
                                      'reasoning_content_start': '<think>'})
    reasoning = Reasoning(model_setting.get('reasoning_content_start', '<think>'),
                          model_setting.get('reasoning_content_end', '</think>'))
    response_reasoning_content = False

    for chunk in response:
        if workflow.is_the_task_interrupted():
            break

        # 处理 reasoning content
        reasoning_chunk = reasoning.get_reasoning_content(chunk)
        content_chunk = reasoning_chunk.get('content')
        if 'reasoning_content' in chunk.additional_kwargs:
            response_reasoning_content = True
            reasoning_content_chunk = chunk.additional_kwargs.get('reasoning_content', '')
        else:
            reasoning_content_chunk = reasoning_chunk.get('reasoning_content')

        answer += content_chunk
        if reasoning_content_chunk is None:
            reasoning_content_chunk = ''
        reasoning_content += reasoning_content_chunk

        # 处理 chunk.content 为 list 的情况
        if isinstance(chunk.content, list):
            for chunk_item in chunk.content:
                text = chunk_item.get("text", "")
                yield {'content': text,
                       'reasoning_content': reasoning_content_chunk if model_setting.get('reasoning_content_enable',
                                                                                         False) else ''}
        else:
            text = chunk.content or ""
            yield {'content': text,
                   'reasoning_content': reasoning_content_chunk if model_setting.get('reasoning_content_enable',
                                                                                     False) else ''}

    reasoning_chunk = reasoning.get_end_reasoning_content()
    answer += reasoning_chunk.get('content')
    reasoning_content_chunk = ""
    if not response_reasoning_content:
        reasoning_content_chunk = reasoning_chunk.get(
            'reasoning_content')
    yield {'content': reasoning_chunk.get('content'),
           'reasoning_content': reasoning_content_chunk if model_setting.get('reasoning_content_enable',
                                                                             False) else ''}
    _write_context(node_variable, workflow_variable, node, workflow, answer, reasoning_content)


def write_context(node_variable: Dict, workflow_variable: Dict, node: INode, workflow):
    """
    写入上下文数据
    @param node_variable:      节点数据
    @param workflow_variable:  全局数据
    @param node:               节点实例对象
    @param workflow:           工作流管理器
    """
    response = node_variable.get('result')
    model_setting = node.context.get('model_setting',
                                     {'reasoning_content_enable': False, 'reasoning_content_end': '</think>',
                                      'reasoning_content_start': '<think>'})
    reasoning = Reasoning(model_setting.get('reasoning_content_start'), model_setting.get('reasoning_content_end'))
    reasoning_result = reasoning.get_reasoning_content(response)
    reasoning_result_end = reasoning.get_end_reasoning_content()
    content = reasoning_result.get('content') + reasoning_result_end.get('content')
    meta = {**response.response_metadata, **response.additional_kwargs}
    if 'reasoning_content' in meta:
        reasoning_content = (meta.get('reasoning_content', '') or '')
    else:
        reasoning_content = (reasoning_result.get('reasoning_content') or '') + (
                reasoning_result_end.get('reasoning_content') or '')
    _write_context(node_variable, workflow_variable, node, workflow, content, reasoning_content)


def file_id_to_base64(file_id: str, video_model):
    file = QuerySet(File).filter(id=file_id).first()
    file_bytes = file.get_bytes()
    url = video_model.upload_file_and_get_url(file_bytes, file.file_name)
    return url


class BaseVideoUnderstandNode(IVideoUnderstandNode):
    def save_context(self, details, workflow_manage):
        self.context['answer'] = details.get('answer')
        self.context['question'] = details.get('question')
        self.context['exception_message'] = details.get('err_message')
        if self.node_params.get('is_result', False):
            self.answer_text = details.get('answer')

    def execute(self, model_id, system, prompt, dialogue_number, dialogue_type, history_chat_record, stream,
                model_params_setting,
                chat_record_id,
                video,
                model_id_type=None, model_id_reference=None,
                model_setting=None,
                **kwargs) -> NodeResult:
        # 处理引用类型
        if model_id_type == 'reference' and model_id_reference:
            reference_data = self.workflow_manage.get_reference_field(
                model_id_reference[0],
                model_id_reference[1:],
            )
            if reference_data and isinstance(reference_data, dict):
                model_id = reference_data.get('model_id', model_id)
                model_params_setting = reference_data.get('model_params_setting')

        from django.utils.translation import gettext_lazy as _

        if model_id is None or model_id == '':
            raise Exception(_('Model is not allowed to be empty'))

        workspace_id = self.workflow_manage.get_body().get('workspace_id')
        if model_setting is None:
            model_setting = {'reasoning_content_enable': False, 'reasoning_content_end': '</think>',
                             'reasoning_content_start': '<think>'}
        self.context['model_setting'] = model_setting
        video_model = get_model_instance_by_model_workspace_id(model_id, workspace_id,
                                                               **(model_params_setting or {}))
        # 执行详情中的历史消息不需要图片内容
        history_message = self.get_history_message_for_details(history_chat_record, dialogue_number)
        self.context['history_message'] = history_message
        question = self.generate_prompt_question(prompt)
        self.context['question'] = question.content
        # 生成消息列表, 真实的history_message
        message_list = self.generate_message_list(video_model, system, prompt,
                                                  self.get_history_message(history_chat_record, dialogue_number,
                                                                           video_model), video)
        self.context['message_list'] = message_list
        self.generate_context_video(video)
        self.context['dialogue_type'] = dialogue_type
        if stream:
            r = video_model.stream(message_list)
            return NodeResult({'result': r, 'chat_model': video_model, 'message_list': message_list,
                               'history_message': history_message, 'question': question.content}, {},
                              _write_context=write_context_stream)
        else:
            r = video_model.invoke(message_list)
            return NodeResult({'result': r, 'chat_model': video_model, 'message_list': message_list,
                               'history_message': history_message, 'question': question.content}, {},
                              _write_context=write_context)

    def generate_context_video(self, video):
        if isinstance(video, str) and video.startswith('http'):
            self.context['video_list'] = [{'url': video}]
        elif video is not None and len(video) > 0:
            self.context['video_list'] = video

    def get_history_message_for_details(self, history_chat_record, dialogue_number):
        start_index = len(history_chat_record) - dialogue_number
        history_message = reduce(lambda x, y: [*x, *y], [
            [self.generate_history_human_message_for_details(history_chat_record[index]),
             self.generate_history_ai_message(history_chat_record[index])]
            for index in
            range(start_index if start_index > 0 else 0, len(history_chat_record))], [])
        return history_message

    def generate_history_ai_message(self, chat_record):
        for val in chat_record.details.values():
            if self.node.id == val['node_id'] and 'video_list' in val:
                if val['dialogue_type'] == 'WORKFLOW':
                    return chat_record.get_ai_message()
                return AIMessage(content=val.get('answer') or val.get('err_message') or '')
        return chat_record.get_ai_message()

    def generate_history_human_message_for_details(self, chat_record):
        for data in chat_record.details.values():
            if self.node.id == data['node_id'] and 'video_list' in data:
                video_list = data['video_list'] or  []
                # 增加对 None 和空列表的检查
                if not video_list or len(video_list) == 0 or data['dialogue_type'] == 'WORKFLOW':
                    return HumanMessage(content=chat_record.problem_text)
                file_id_list = []
                url_list = []
                for image in video_list:
                    if 'file_id' in image:
                        file_id_list.append(image.get('file_id'))
                    elif 'url' in image:
                        url_list.append(image.get('url'))
                return HumanMessage(content=[
                    {'type': 'text', 'text': data['question']},
                    *[{'type': 'video_url', 'video_url': {'url': f'./oss/file/{file_id}'}} for file_id in file_id_list],
                    *[{'type': 'video_url', 'video_url': {'url': url}} for url in url_list],
                ])
        return HumanMessage(content=chat_record.problem_text)

    def get_history_message(self, history_chat_record, dialogue_number, video_model):
        start_index = len(history_chat_record) - dialogue_number
        history_message = reduce(lambda x, y: [*x, *y], [
            [self.generate_history_human_message(history_chat_record[index], video_model),
             self.generate_history_ai_message(history_chat_record[index])]
            for index in
            range(start_index if start_index > 0 else 0, len(history_chat_record))], [])
        return history_message

    def generate_history_human_message(self, chat_record, video_model):

        for data in chat_record.details.values():
            if self.node.id == data['node_id'] and 'video_list' in data:
                video_list = data['video_list'] or  []
                if video_list is None or len(video_list) == 0 or data['dialogue_type'] == 'WORKFLOW':
                    return HumanMessage(content=chat_record.problem_text)
                file_id_list = []
                url_list = []
                for image in video_list:
                    if 'file_id' in image:
                        file_id_list.append(image.get('file_id'))
                    elif 'url' in image:
                        url_list.append(image.get('url'))
                video_base64_list = [file_id_to_base64(video.get('file_id'), video_model) for video in video_list]
                return HumanMessage(
                    content=[
                        {'type': 'text', 'text': data['question']},
                        *[{'type': 'video_url',
                           'video_url': {'url': f'{base64_video}'}} for
                          base64_video in video_base64_list]
                    ])
        return HumanMessage(content=chat_record.problem_text)

    def generate_prompt_question(self, prompt):
        return HumanMessage(self.workflow_manage.generate_prompt(prompt))

    def _process_videos(self, image, video_model):
        videos = []
        if isinstance(image, str) and image.startswith('http'):
            videos.append({'type': 'video_url', 'video_url': {'url': image}})
        elif image is not None and len(image) > 0:
            for img in image:
                if 'file_id' in img:
                    file_id = img['file_id']
                    file = QuerySet(File).filter(id=file_id).first()
                    url = video_model.upload_file_and_get_url(file.get_bytes(), file.file_name)
                    videos.append(
                        {'type': 'video_url', 'video_url': {'url': url}})
                elif 'url' in img and img['url'].startswith('http'):
                    videos.append(
                        {'type': 'video_url', 'video_url': {'url': img['url']}})
        return videos

    def generate_message_list(self, video_model, system: str, prompt: str, history_message, video):
        prompt_text = self.workflow_manage.generate_prompt(prompt)
        videos = self._process_videos(video, video_model)

        if videos:
            messages = [HumanMessage(content=[{'type': 'text', 'text': prompt_text}, *videos])]
        else:
            messages = [HumanMessage(prompt_text)]

        if system is not None and len(system) > 0:
            return [
                SystemMessage(self.workflow_manage.generate_prompt(system)),
                *history_message,
                *messages
            ]
        else:
            return [
                *history_message,
                *messages
            ]

    @staticmethod
    def reset_message_list(message_list: List[BaseMessage], answer_text):
        result = [{'role': 'user' if isinstance(message, HumanMessage) else 'ai', 'content': message.content} for
                  message
                  in
                  message_list]
        result.append({'role': 'ai', 'content': answer_text})
        return result

    def get_details(self, index: int, **kwargs):
        return {
            'name': self.node.properties.get('stepName'),
            "index": index,
            'run_time': self.context.get('run_time'),
            'system': self.node_params.get('system'),
            'history_message': [{'content': message.content, 'role': message.type} for message in
                                (self.context.get('history_message') if self.context.get(
                                    'history_message') is not None else [])],
            'question': self.context.get('question'),
            'answer': self.context.get('answer'),
            'reasoning_content': self.context.get('reasoning_content'),
            'type': self.node.type,
            'message_tokens': self.context.get('message_tokens'),
            'answer_tokens': self.context.get('answer_tokens'),
            'status': self.status,
            'err_message': self.err_message,
            'video_list': self.context.get('video_list'),
            'dialogue_type': self.context.get('dialogue_type'),
            'enableException': self.node.properties.get('enableException'),
        }
