import torch
import random
from tqdm import tqdm
from comfy.model_base import BaseModel
from comfy.sd import load_checkpoint_guess_config
from server import PromptServer
import os
import subprocess
import sys
import requests
import hashlib
import folder_paths

def ensure_openai():
    if "python_embeded" in sys.executable or "python_embedded" in sys.executable:
        pip_install = [sys.executable, "-s", "-m", "pip", "install"]
    else:
        pip_install = [sys.executable, "-m", "pip", "install"]

    try:
        from openai import OpenAI
    except Exception as e:
        try:
            subprocess.check_call(pip_install + ['openai'])
            from openai import OpenAI
        except:
            print("Failed to install 'openai'. Please, install manually.")

from openai import OpenAI

import json


class FileUploader:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "file_path": ("STRING", {"default": ""}),
                "purpose": (["file-extract", "retrieval"], {"default": "retrieval"}),
                "api_key": ("STRING", {"default": ""})
            },
            "hidden": {
                "unique_id": "UNIQUE_ID"
            }
        }
    
    RETURN_TYPES = ("STRING", "STRING", "STRING", "STRING")
    RETURN_NAMES = ("file_id", "filename", "purpose", "status")
    FUNCTION = "upload_file"
    CATEGORY = "StepFun"

    def upload_file(self, file_path, purpose, api_key, unique_id):
        if not os.path.exists(file_path):
            raise ValueError(f"文件不存在: {file_path}")
            
        headers = {
            "Authorization": f"Bearer {api_key}"
        }
        
        PromptServer.instance.send_sync("stepfun.file.upload", {
            "node": unique_id,
            "message": f"开始上传文件: {os.path.basename(file_path)}"
        })
        
        with open(file_path, "rb") as f:
            files = {
                "file": f,
                "purpose": (None, purpose)
            }
            
            response = requests.post(
                "https://api.stepfun.com/v1/files",
                headers=headers,
                files=files
            )
            
        if response.status_code != 200:
            error_msg = f"上传失败: {response.text}"
            PromptServer.instance.send_sync("stepfun.file.upload", {
                "node": unique_id,
                "message": error_msg,
                "type": "error"
            })
            raise Exception(error_msg)
            
        result = response.json()
        
        PromptServer.instance.send_sync("stepfun.file.upload", {
            "node": unique_id,
            "message": f"文件上传成功: {result.get('filename', '')}",
            "type": "success"
        })
        
        return {
            "ui": {
                "status": result.get("status", ""),
                "filename": result.get("filename", "")
            },
            "result": (
                result.get("id", ""),
                result.get("filename", ""),
                result.get("purpose", ""),
                result.get("status", "")
            )
        }

class StepFunClient:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "api_key": ("STRING", {"default": ""})
            }
        }
    
    RETURN_TYPES = ("STEPFUN_CLIENT",)
    RETURN_NAMES = ("client",)
    FUNCTION = "create_client"
    CATEGORY = "StepFun"

    def create_client(self, api_key):
        client = OpenAI(
            base_url="https://api.stepfun.com/v1",
            api_key=api_key
        )
        return (client,)

class TextImageChat:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "client": ("STEPFUN_CLIENT",),
                "model": (["step-1v-8k", "step-1v-32k", "step-1.5v-mini","step-2-16k"], {"default": "step-1.5v-mini"}),
                "detail": (["low", "high"], {"default": "low"}),
                "system_prompt": ("STRING", {
                    "default": "Describe the picture",
                    "multiline": True
                }),
                "user_prompt": ("STRING", {
                    "default": "",
                    "multiline": True
                }),
                "json_mode": ("BOOLEAN", {"default": False}),
                "temperature": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 2.0, "step": 0.1}),
                "top_p": ("FLOAT", {"default": 0.9, "min": 0.0, "max": 1.0, "step": 0.1}),
                "max_tokens": ("INT", {"default": 2048, "min": 1, "max": 8192})
            },           
            "optional": {
                "image": ("IMAGE",),
            },
        }
    
    RETURN_TYPES = ("STRING", "STRING", "STRING")
    RETURN_NAMES = ("response", "total_tokens", "finish_reason")
    FUNCTION = "chat_completion"
    CATEGORY = "StepFun"
    OUTPUT_NODE = True
    def process_image(self, image, detail="low"):
        """处理图片尺寸和质量"""
        from PIL import Image
        import numpy as np
        import io
        import base64
        
        i = 255. * image[0].cpu().numpy()
        img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
        
        max_size = 2688 if detail == "high" else 1280
        width, height = img.size
        
        if width > height:
            new_width = max_size
            new_height = int((max_size / width) * height)
        else:
            new_height = max_size
            new_width = int((max_size / height) * width)
        
        img = img.resize((new_width, new_height), Image.LANCZOS)
        
        buffer = io.BytesIO()
        img.save(buffer, format="PNG", quality=80)
        img_str = base64.b64encode(buffer.getvalue()).decode()
        
        return img_str

    def chat_completion(self, client, model, system_prompt, user_prompt, json_mode, temperature, top_p, max_tokens, detail, image=None):
        messages = [
            {
                "role": "system",
                "content": system_prompt
            }
        ]
        
        user_message = {
            "role": "user",
            "content": []
        }
        
        if user_prompt:
            user_message["content"].append({
                "type": "text",
                "text": user_prompt
            })
        
        if image is not None and not model.startswith("step-2"):
            img_str = self.process_image(image, detail)
            user_message["content"].append({
                "type": "image_url",
                "image_url": {
                    "url": f"data:image/png;base64,{img_str}",
                    "detail": detail
                }
            })
        
        messages.append(user_message)
        
        kwargs = {
            "model": model,
            "messages": messages,
            "temperature": temperature,
            "top_p": top_p,
            "max_tokens": max_tokens
        }
        
        if json_mode:
            kwargs["response_format"] = {"type": "json_object"}
        
        completion = client.chat.completions.create(**kwargs)
        
        response_content = completion.choices[0].message.content
        total_tokens = str(completion.usage.total_tokens)
        finish_reason = completion.choices[0].finish_reason
        
        return {
            "ui": {
                "response": response_content
            },
            "result": (response_content, total_tokens, finish_reason)
        }

class VideoChat:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "client": ("STEPFUN_CLIENT",),
                "model": (["step-1v-8k", "step-1v-32k", "step-1.5v-mini"], {"default": "step-1.5v-mini"}),
                "system_prompt": ("STRING", {
                    "default": "Description What happened to the uploaded video",
                    "multiline": True
                }),
                "user_prompt": ("STRING", {
                    "default": "",
                    "multiline": True
                }),
                "video_url": ("STRING", {"default": ""}),
                "json_mode": ("BOOLEAN", {"default": False}),
                "temperature": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 2.0, "step": 0.1}),
                "top_p": ("FLOAT", {"default": 0.9, "min": 0.0, "max": 1.0, "step": 0.1}),
                "max_tokens": ("INT", {"default": 1024, "min": 1, "max": 8192})
            },
            "hidden": {
                "unique_id": "UNIQUE_ID"
            }
        }
    
    RETURN_TYPES = ("STRING", "STRING", "STRING")
    RETURN_NAMES = ("response", "total_tokens", "finish_reason")
    FUNCTION = "chat_completion"
    CATEGORY = "StepFun"
    OUTPUT_NODE = True

    def chat_completion(self, client, model, system_prompt, user_prompt, video_url, json_mode, temperature, top_p, max_tokens, unique_id):
        messages = [
            {
                "role": "system",
                "content": system_prompt
            }
        ]
        
        user_message = {
            "role": "user",
            "content": []
        }
        
        if user_prompt:
            user_message["content"].append({
                "type": "text",
                "text": user_prompt
            })
        
        if video_url:
            user_message["content"].append({
                "type": "video_url",
                "video_url": {
                    "url": video_url
                }
            })
        
        messages.append(user_message)
        
        kwargs = {
            "model": model,
            "messages": messages,
            "temperature": temperature,
            "top_p": top_p,
            "max_tokens": max_tokens
        }
        
        if json_mode:
            kwargs["response_format"] = {"type": "json_object"}
        
        completion = client.chat.completions.create(**kwargs)
        
        response_content = completion.choices[0].message.content
        total_tokens = str(completion.usage.total_tokens)
        finish_reason = completion.choices[0].finish_reason
        
        return {
            "ui": {
                "response": response_content
            },
            "result": (response_content, total_tokens, finish_reason)
        }

class JSONParser:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "json_string": ("STRING", {"multiline": True}),
            }
        }
    
    RETURN_TYPES = ()  # 动态生成
    RETURN_NAMES = ()  # 动态生成
    FUNCTION = "parse_json"
    CATEGORY = "StepFun"

    def parse_json(self, json_string):
        try:
            data = json.loads(json_string)
            if not isinstance(data, dict):
                raise ValueError("JSON必须是一个对象")
                
            # 动态生成输出类型和名称
            self.__class__.RETURN_TYPES = ("STRING",) * len(data)
            self.__class__.RETURN_NAMES = tuple(data.keys())
            
            # 将所有值转换为字符串
            values = []
            for key in data.keys():
                val = data[key]
                if isinstance(val, (dict, list)):
                    values.append(json.dumps(val, ensure_ascii=False))
                else:
                    values.append(str(val))
                    
            return tuple(values)
            
        except Exception as e:
            raise ValueError(f"JSON解析错误: {str(e)}")


class CombineStrings:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "input_1": ("STRING", {"multiline": True}),
                "input_2": ("STRING", {"multiline": True}),
                "input_3": ("STRING", {"multiline": True})
            }
        }
    
    RETURN_TYPES = ("STRING_LIST", "STRING")
    RETURN_NAMES = ("string_array", "combined_string")
    FUNCTION = "combine"
    CATEGORY = "StepFun"

    def combine(self, input_1, input_2, input_3):
        # 收集所有非空输入的字符串
        all_strings = [s.strip() for s in [input_1, input_2, input_3] if s.strip()]
        
        # 用逗号连接
        combined_string = ','.join(all_strings)
        
        return (all_strings, combined_string)

class VideoFileUploader:
    # 用于存储上传记录的缓存文件路径
    CACHE_FILE = os.path.join(os.path.dirname(__file__), "video_upload_cache.json")
    
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "client": ("STEPFUN_CLIENT",),
                "file_path": ("STRING", {"default": ""}),
                "force_reupload": ("BOOLEAN", {"default": False}),
            },
            "hidden": {
                "unique_id": "UNIQUE_ID"
            }
        }
    
    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("file_link",)
    FUNCTION = "upload_video_file"
    CATEGORY = "StepFun"

    def _load_cache(self):
        """加载缓存数据"""
        if os.path.exists(self.CACHE_FILE):
            try:
                with open(self.CACHE_FILE, 'r') as f:
                    return json.load(f)
            except:
                return {}
        return {}

    def _save_cache(self, cache_data):
        """保存缓存数据"""
        os.makedirs(os.path.dirname(self.CACHE_FILE), exist_ok=True)
        with open(self.CACHE_FILE, 'w') as f:
            json.dump(cache_data, f)

    def _get_file_hash(self, file_path):
        """计算文件的MD5哈希值"""
        hash_md5 = hashlib.md5()
        with open(file_path, "rb") as f:
            for chunk in iter(lambda: f.read(4096), b""):
                hash_md5.update(chunk)
        return hash_md5.hexdigest()

    def upload_video_file(self, client, file_path, force_reupload, unique_id):
        if not os.path.exists(file_path):
            raise ValueError(f"文件不存在: {file_path}")
        
        # 加载缓存
        cache = self._load_cache()
        file_hash = self._get_file_hash(file_path)
        
        # 检查缓存中是否存在且不强制重新上传
        if not force_reupload and file_hash in cache:
            PromptServer.instance.send_sync("stepfun.file.upload", {
                "node": unique_id,
                "message": f"使用缓存的上传记录: {os.path.basename(file_path)}",
                "type": "success"
            })
            return {
                "ui": {
                    "file_link": cache[file_hash]
                },
                "result": (cache[file_hash],)
            }
        
        PromptServer.instance.send_sync("stepfun.file.upload", {
            "node": unique_id,
            "message": f"开始上传视频文件: {os.path.basename(file_path)}"
        })
        
        try:
            response = client.files.create(
                file=open(file_path, "rb"),
                purpose="storage"
            )
            
            file_link = f"stepfile://{response.id}"
            
            # 更新缓存
            cache[file_hash] = file_link
            self._save_cache(cache)
            
            PromptServer.instance.send_sync("stepfun.file.upload", {
                "node": unique_id,
                "message": f"视频文件上传成功: {response.filename}",
                "type": "success"
            })
            
            return {
                "ui": {
                    "file_link": file_link
                },
                "result": (file_link,)
            }
            
        except Exception as e:
            error_msg = f"上传失败: {str(e)}"
            PromptServer.instance.send_sync("stepfun.file.upload", {
                "node": unique_id,
                "message": error_msg,
                "type": "error"
            })
            raise Exception(error_msg)

# 更新节点映射
NODE_CLASS_MAPPINGS = {
    "StepFunClient": StepFunClient,
    "TextImageChat": TextImageChat,
    "VideoChat": VideoChat,
    "JSONParser": JSONParser,
    "CombineStrings": CombineStrings,
    "VideoFileUploader": VideoFileUploader,

}

NODE_DISPLAY_NAME_MAPPINGS = {
    "StepFunClient": "StepFun Client",
    "TextImageChat": "StepFun Chat Completion", 
    "VideoChat": "StepFun Video Chat",
    "JSONParser": "StepFun JSON Parser",
    "CombineStrings": "StepFun String Combiner",
    "VideoFileUploader": "StepFun Video File Uploader",
    
}