import json
import requests
import logging

from JiMeng.Kit.VolcanoConst import VOLCANO_API_KEY

# 配置日志
logger = logging.getLogger(__name__)
logger.level = logging.INFO
console_handler = logging.StreamHandler()
console_handler.setLevel(logging.INFO)
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
console_handler.setFormatter(formatter)
logger.addHandler(console_handler)

class FenJingTouGenerator:
    def __init__(self, api_key=VOLCANO_API_KEY):
        """
        初始化分镜头提示词生成器
        :param api_key: 火山引擎API密钥
        """
        self.api_key = api_key
        self.url = 'https://ark.cn-beijing.volces.com/api/v3/chat/completions'
        self.headers = {
            'Content-Type': 'application/json',
            'Authorization': f'Bearer {self.api_key}'
        }
        self.model = 'doubao-1-5-vision-pro-32k-250115'

    def generate_fen_jing_tou(self, image_url, prompt_text, timeout=60):
        """
        生成分镜头提示词
        :param image_url: 首帧图片URL
        :param prompt_text: 提示文本
        :param timeout: 请求超时时间(秒)
        :return: 生成的分镜头提示词内容，如果失败则返回None
        """
        # 构建请求体
        payload = {
            'model': self.model,
            'messages': [
                {
                    'content': [
                        {
                            'image_url': {
                                'url': image_url
                            },
                            'type': 'image_url'
                        },
                        {
                            'text': prompt_text,
                            'type': 'text'
                        }
                    ],
                    'role': 'user'
                }
            ]
        }

        try:
            # 发送POST请求
            response = requests.post(
                self.url,
                headers=self.headers,
                data=json.dumps(payload),
                timeout=timeout
            )
            # 检查响应状态码
            response.raise_for_status()
            # 解析响应数据
            result = response.json()

            # 提取生成的内容
            if 'choices' in result and len(result['choices']) > 0:
                choice = result['choices'][0]
                if 'message' in choice and 'content' in choice['message']:
                    content = choice['message']['content']
                    logger.info(f'分镜头提示词生成成功')
                    return content
                else:
                    logger.error('响应中未找到message或content字段')
                    return None
            else:
                logger.error('响应中未找到choices字段或choices为空')
                return None

        except requests.exceptions.RequestException as e:
            logger.error(f'请求发生错误: {e}')
            if hasattr(e, 'response') and e.response is not None:
                logger.error(f'错误响应状态码: {e.response.status_code}')
                logger.error(f'错误响应内容: {e.response.text}')
            return None

        except json.JSONDecodeError:
            logger.error('响应解析失败: 返回的内容不是有效的JSON格式')
            if 'response' in locals():
                logger.error(f'响应内容: {response.text}')
            return None




if __name__ == '__main__':
    # 创建实例并运行
    generator = FenJingTouGenerator()

    """主函数，用于测试分镜头提示词生成功能"""
    image_url = 'https://dsideal.obs.cn-north-1.myhuaweicloud.com/HuangHai/JiMeng/d4186502-5d6e-404b-890c-1f40fbbc1728.jpg'
    prompt_text = '以这张图片为首帧，帮我生成一个5秒的视频提示词。提示词包括：镜号、运镜、画面内容'

    # 生成分镜头提示词
    content = generator.generate_fen_jing_tou(image_url, prompt_text)

    # 打印结果
    if content:
        print(content)
    else:
        print('分镜头提示词生成失败')