import os
import random
import re
import json
import urllib.request
import urllib.error
import base64
import io
import time
import numpy as np
from PIL import Image
import torch


class DolphinMultiActionPromptNode_V32:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "image": ("IMAGE",),
                "mode": (["🤖 Auto Vision+LLM", "✍️ Manual Override"], {"default": "🤖 Auto Vision+LLM"}),
                "character_name": ("STRING", {"multiline": False, "default": "AUTO"}),
                "artistic_vibe": ("STRING", {"multiline": True, "default": "cinematic lighting, high-speed action, dark fantasy"}),

                "master_story": ("STRING", {
                    "multiline": True,
                    "default": "어두운 골목길. 갑자기 나타난 적들을 향해 돌진한다, 화려하게 검을 휘둘러 적을 쓰러뜨린다, 날아오는 총알을 튕겨낸다, 적에게 다가가 숨통을 끊는다."
                }),

                "openrouter_api_key": ("STRING", {"multiline": False, "default": ""}),
                "openrouter_model": ("STRING", {"multiline": False, "default": "qwen/qwen-2-vl-72b-instruct"}),
                "creativity": ("FLOAT", {"default": 0.85, "min": 0.1, "max": 1.5, "step": 0.05}),
                "max_tokens": ("INT", {"default": 1500, "min": 256, "max": 8192, "step": 64}),
                "retries": ("INT", {"default": 2, "min": 0, "max": 5}),
                "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
            },
        }

    RETURN_TYPES = ("STRING", "STRING", "STRING", "STRING", "STRING")
    RETURN_NAMES = (
        "prompt_1 (Clip 1: 0-5s)",
        "prompt_2 (Clip 2: 0-5s)",
        "prompt_3 (Clip 3: 0-5s)",
        "prompt_4 (Clip 4: 0-5s)",
        "raw_llm_output",
    )
    FUNCTION = "generate_sequence"
    CATEGORY = "Dolphin"

    # -----------------------------------------------------------------
    def _encode_image(self, image):
        img_tensor = image[0]
        i = 255. * img_tensor.cpu().numpy()
        img_pil = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
        if img_pil.mode != "RGB":
            img_pil = img_pil.convert("RGB")
        buffered = io.BytesIO()
        img_pil.save(buffered, format="JPEG", quality=90)
        b64 = base64.b64encode(buffered.getvalue()).decode('utf-8')
        return f"data:image/jpeg;base64,{b64}"

    def _split_sentences(self, master_story):
        clean_story = re.sub(r'([.!?,\n])', r'\1|', master_story)
        raw_sentences = clean_story.split('|')
        return [s.strip() for s in raw_sentences if len(s.strip()) > 1]

    def _build_chunks(self, sentences):
        n = len(sentences)
        if n == 0:
            base = "dynamic high-speed action"
            return (base, base, base, base)
        if n == 1:
            s = sentences[0]
            return (
                f"Phase 1: Rapid approach and high-speed dynamic movement. DO NOT stand still. (Target: {s})",
                f"Phase 2: Swift, explosive execution of the action. (Target: {s})",
                f"Phase 3: The climax at full 1x real-time speed. Lightning fast! (Target: {s})",
                f"Phase 4: Fast-paced completion and quick recovery. (Target: {s})",
            )
        if n == 2:
            return (
                f"Phase 1: High-speed buildup and rapid preparation. (Target: {sentences[0]})",
                f"Phase 2: Explosively execute -> {sentences[0]}",
                f"Phase 3: Rapid transition, sprinting or moving quickly. (Target: {sentences[1]})",
                f"Phase 4: Lightning-fast execution -> {sentences[1]}",
            )
        if n == 3:
            return (
                f"Phase 1: Start this action rapidly -> {sentences[0]}",
                f"Phase 2: Explosively complete -> {sentences[0]}",
                sentences[1],
                sentences[2],
            )
        # n >= 4: 균등 분배
        k, m = divmod(n, 4)
        chunks = []
        start = 0
        for idx in range(4):
            end = start + k + (1 if idx < m else 0)
            chunks.append(" ".join(sentences[start:end]))
            start = end
        return tuple(chunks)

    def _call_llm(self, url, payload, api_key, retries, timeout=120):
        last_err = None
        for attempt in range(retries + 1):
            try:
                req = urllib.request.Request(
                    url,
                    data=json.dumps(payload).encode('utf-8'),
                    headers={'Authorization': f'Bearer {api_key}', 'Content-Type': 'application/json'}
                )
                response = urllib.request.urlopen(req, timeout=timeout)
                body = json.loads(response.read().decode('utf-8'))
                return body['choices'][0]['message']['content'].strip(), None
            except Exception as e:
                last_err = e
                if attempt < retries:
                    time.sleep(1.5 * (attempt + 1))
        return None, last_err

    # -----------------------------------------------------------------
    def generate_sequence(self, image, mode, character_name, artistic_vibe, master_story,
                          openrouter_api_key, openrouter_model, creativity, max_tokens, retries, seed):

        user_defined_name = "" if character_name.upper() in ["AUTO", ""] else character_name.strip()

        def build_final(action, master_scene, name):
            tags_list = []
            if name:
                tags_list.append(name)
            tag_block = ", ".join(tags_list)

            sentence_list = []
            if master_scene.strip():
                sentence_list.append(master_scene.strip().strip(",. "))
            if action.strip():
                sentence_list.append(action.strip())
            sentence_block = " ".join(sentence_list)

            if tag_block and sentence_block:
                return f"{tag_block}\n{sentence_block}"
            elif tag_block:
                return tag_block
            return sentence_block

        # ---- Manual Override ----
        if mode == "✍️ Manual Override":
            fp = build_final(master_story, "", user_defined_name)
            return (fp, fp, fp, fp, "[Manual Override]")

        random.seed(seed)

        base64_image = self._encode_image(image)
        sentences = self._split_sentences(master_story)
        chunk_1, chunk_2, chunk_3, chunk_4 = self._build_chunks(sentences)

        sys_prompt = (
            "You are an Elite Action Director prioritizing RAW SPEED and KINETIC ENERGY.\n"
            f"1. CHARACTER: If NAME is 'AUTO', assign a name. If '{user_defined_name}', use it.\n"
            "2. VISUAL ANALYSIS: You MUST base your descriptions EXACTLY on the character's clothing and weapons in the attached IMAGE.\n"
            "3. MASTER SCENE: Write a 1-sentence environment description (lighting, weather).\n"
            "4. SPEED-FOCUSED CHOREOGRAPHY (CRITICAL):\n"
            "   - 🚫 BAN SLOW-MOTION TRIGGERS: NEVER use words like 'micro-expressions', 'muscle tension', 'slowly turning', 'floating', or 'gradually'. These cause AI video models to render in slow-motion.\n"
            "   - ✅ FORCE 1x REAL-TIME SPEED: Describe large, sweeping, high-velocity movements. Use aggressive verbs (dashing, sprinting, whipping, snapping).\n"
            "   - ✅ KINETIC ADVERBS: Inject phrases like 'in a flash', 'at lightning speed', 'with explosive real-time velocity' into EVERY part.\n"
            "   - Example: 'suddenly dashes forward at full speed and delivers a lightning-fast horizontal strike, moving so quickly the rain splatters'.\n"
            "   - Strictly confine the actions. DO NOT animate future events early.\n"
            "   - 🔥 OUTPUT RULE: DO NOT quote the Korean text. Only output English.\n"
            "Format EXACTLY:\nCHARACTER: [Name]\nMASTER SCENE: [Description]\n"
            "PART 1: [0-1s] [Action A] [2-3s] [Action B] [4-5s] [Action C]\n"
            "PART 2: [0-1s] [Action D] [2-3s] [Action E] [4-5s] [Action F]\n"
            "PART 3: [0-1s] [Action G] [2-3s] [Action H] [4-5s] [Action I]\n"
            "PART 4: [0-1s] [Action J] [2-3s] [Action K] [4-5s] [Action L]"
        )

        usr_text = (
            f"NAME: {character_name}\n"
            f"VIBE: {artistic_vibe}\n\n"
            "=== HIGH-SPEED ACTION SCRIPT ===\n"
            f"▶ For PART 1 (0-5s), ONLY animate this: \"{chunk_1}\"\n"
            f"▶ For PART 2 (5-10s), ONLY animate this: \"{chunk_2}\"\n"
            f"▶ For PART 3 (10-15s), ONLY animate this: \"{chunk_3}\"\n"
            f"▶ For PART 4 (15-20s), ONLY animate this: \"{chunk_4}\"\n"
            "CRITICAL: Keep the action moving FAST. Avoid still poses or micro-details that look like slow-mo!"
        )

        url = "https://openrouter.ai/api/v1/chat/completions"
        payload = {
            "model": openrouter_model.strip(),
            "messages": [
                {"role": "system", "content": sys_prompt},
                {"role": "user", "content": [
                    {"type": "text", "text": usr_text},
                    {"type": "image_url", "image_url": {"url": base64_image}}
                ]}
            ],
            "temperature": creativity,
            "max_tokens": max_tokens,
        }

        if not openrouter_api_key.strip():
            err_msg = "⚠️ API Error: OpenRouter API key is empty."
            return (err_msg, err_msg, err_msg, err_msg, err_msg)

        llm_prompt, err = self._call_llm(url, payload, openrouter_api_key, retries)
        if llm_prompt:
            print(f"\n✅ [Dolphin V32 - Action Speed Optimized]\n{llm_prompt}\n")
        else:
            print(f"❌ [에러] API 호출 실패: {err}")
            llm_prompt = ""

        final_char_name = user_defined_name
        r_master = ""
        p1 = p2 = p3 = p4 = ""

        if llm_prompt and "[removed]" not in llm_prompt:
            cl = re.sub(r'[*#]', '', llm_prompt)
            m_char = re.search(r'CHARACTER:\s*(.*?)(?=MASTER SCENE|$)', cl, re.I | re.S)
            m_master = re.search(r'MASTER SCENE:\s*(.*?)(?=PART 1|$)', cl, re.I | re.S)

            m1 = re.search(r'PART 1:\s*(.*?)(?=PART 2|$)', cl, re.I | re.S)
            m2 = re.search(r'PART 2:\s*(.*?)(?=PART 3|$)', cl, re.I | re.S)
            m3 = re.search(r'PART 3:\s*(.*?)(?=PART 4|$)', cl, re.I | re.S)
            m4 = re.search(r'PART 4:\s*(.*?)(?=\n\n|===|Note:|$)', cl, re.I | re.S)

            if not user_defined_name and m_char:
                final_char_name = m_char.group(1).strip()

            r_master = m_master.group(1).strip() if m_master else ""

            # 파싱 실패 시 해당 청크(영문 지시문)를 폴백으로 사용해 비디오 프롬프트가 비지 않도록 함
            p1 = m1.group(1).strip() if m1 else chunk_1
            p2 = m2.group(1).strip() if m2 else chunk_2
            p3 = m3.group(1).strip() if m3 else chunk_3
            p4 = m4.group(1).strip() if m4 else chunk_4
        else:
            # API 실패 시에도 스크립트 청크를 폴백으로 반환 (완전 실패보다 유용)
            p1, p2, p3, p4 = chunk_1, chunk_2, chunk_3, chunk_4

        return (
            build_final(p1, r_master, final_char_name),
            build_final(p2, r_master, final_char_name),
            build_final(p3, r_master, final_char_name),
            build_final(p4, r_master, final_char_name),
            llm_prompt if llm_prompt else "⚠️ API Error / empty response",
        )