import asyncio
import logging
from typing import List, Dict, Any, Optional
import os

from tenacity import retry, stop_after_attempt, wait_exponential

from agentscope.message import Msg
from agentscope.agent import ReActAgent
from agentscope.formatter import OpenAIChatFormatter

from src.services.provider.registry import ModelRegistry, ModelType
from src.services.provider.agentscope_adapter import PixelAgentScopeModel
from ..models import (
    CharacterList, SceneList, PropList, NovelSummary,
    DeduplicationResult, StoryboardList, ScriptDiscussion, ScreenplayDraft
)
from ..prompts import (
    SUMMARIZER_PROMPT,
    CHARACTER_EXTRACTOR_PROMPT, SCENE_EXTRACTOR_PROMPT, PROP_EXTRACTOR_PROMPT,
    DEDUPLICATOR_PROMPT, STORYBOARD_ARTIST_PROMPT,
    SCRIPT_DISCUSSION_PROMPT, SCREENPLAY_ADAPTATION_PROMPT,
    OPTIMIZATION_PROMPT_IMAGE_DEFAULT, OPTIMIZATION_PROMPT_VIDEO_DEFAULT,
    OPTIMIZATION_PROMPT_DEFAULT, OPTIMIZATION_PROMPT_IMAGE_STORYBOARD,
    OPTIMIZATION_PROMPT_VIDEO_STORYBOARD_INTEGRATED, OPTIMIZATION_PROMPT_IMAGE_WHITE_BG,
    OPTIMIZATION_PROMPT_IMAGE_THREE_VIEW, OPTIMIZATION_PROMPT_VIDEO_ASSET_360
)
from ..utils import build_context, clean_json_string
import json
import json_repair

from src.constants.common import CHARACTER_ROLES, CHARACTER_GENDERS, SHOT_TYPES, CAMERA_MOVEMENTS, TRANSITIONS

from .config import GENERATION_OPTIONS

logger = logging.getLogger(__name__)

# Use options from config module
_OPTIONS = GENERATION_OPTIONS

class ScriptAnalysisPipeline:
    def __init__(self, model_name: Optional[str] = None, base_url: Optional[str] = None, api_key: Optional[str] = None, provider: Optional[str] = None):
        self.model_name = model_name or "qwen-plus"
        self.provider = provider

        # Resolve provider and model_name for PixelAgentScopeModel(provider, model_name)
        if "/" in self.model_name:
            self.provider, self.model_name = self.model_name.split("/", 1)
        else:
            self.provider = self.provider or "dashscope"

        # If provider was specified, try to resolve service id from registry
        if provider and "/" not in (model_name or ""):
            try:
                services = ModelRegistry.find_services(provider=provider, model_type=ModelType.LLM)
                for svc in services:
                    if svc.get("model_name") == self.model_name or svc.get("id") == self.model_name:
                        rid = svc.get("id", "")
                        if "/" in str(rid):
                            self.provider, self.model_name = str(rid).split("/", 1)
                        logger.info(f"Resolved service config for {provider}/{self.model_name}")
                        break
            except Exception as e:
                logger.warning(f"Failed to resolve service from registry: {e}")

        self._base_url = base_url
        self._api_key = api_key
        # 延迟初始化 model_config，避免启动时需要 API Key
        self._model_config = None

    @property
    def model_config(self):
        """延迟初始化模型配置"""
        if self._model_config is None:
            overrides = {}
            if self._base_url:
                overrides["base_url"] = self._base_url
            if self._api_key:
                overrides["api_key"] = self._api_key

            self._model_config = PixelAgentScopeModel(
                provider=self.provider,
                model_name=self.model_name,
                **overrides
            )
        return self._model_config

    def set_api_key(self, api_key: str):
        """动态设置 API Key（用于纯用户密钥模式）"""
        self._api_key = api_key
        if self._model_config is not None:
            self._model_config.set_api_key(api_key)

    def _create_agent(self, name: str, sys_prompt: str) -> ReActAgent:
        return ReActAgent(
            name=name,
            sys_prompt=sys_prompt,
            model=self.model_config,
            formatter=OpenAIChatFormatter(),
        )

    @retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=2, max=10))
    async def _call_agent_with_retry(self, agent, msg, structured_model=None):
        return await agent(msg, structured_model=structured_model)

    def _salvage_novel_summary_from_response(self, response: Msg) -> tuple:
        """从 response 的 content 中尝试提取标题和摘要（用于 generate_response 参数为空时的兜底）。"""
        content = response.content
        if isinstance(content, list):
            text_content = ""
            for item in content:
                if isinstance(item, dict) and item.get("type") == "text":
                    text_content += item.get("text", "")
            content = text_content
        if not content or not isinstance(content, str):
            return "", ""
        content = content.strip()
        if not content:
            return "", ""
        lines = [ln.strip() for ln in content.split("\n") if ln.strip()]
        if not lines:
            return "未命名", content
        # 第一行作为标题，其余作为摘要；若只有一段则整段作为摘要
        if len(lines) == 1:
            return "未命名", lines[0]
        return lines[0], "\n".join(lines[1:])

    async def _parse_agent_response(self, response: Msg, response_model: Any) -> Any:
        """ 辅助函数 to parse structured output from agent response."""
        # Check if metadata contains the parsed object
        if hasattr(response, 'metadata') and response.metadata:
            if hasattr(response.metadata, "model_dump"):
                return response.metadata.model_dump()
            return response.metadata

        # Fallback: Parse content manually
        content = response.content
        if isinstance(content, list):
            text_content = ""
            for item in content:
                if isinstance(item, dict) and item.get("type") == "text":
                    text_content += item.get("text", "")
            content = text_content

        try:
            return json_repair.loads(content)
        except Exception as e:
            logger.warning(f"json_repair failed, trying standard cleanup: {e}")
            cleaned = clean_json_string(content)
            return json.loads(cleaned)

    def _normalize_extraction_result(self, result: Any, list_key: str) -> Dict[str, Any]:
        """Normalize agent response to dict with list_key, is_complete, continuation_context.
        Handles: dict, list (treat as list of items), or Pydantic model.
        """
        def ensure_dict_list(items: List) -> List[Dict]:
            out = []
            for x in items or []:
                if hasattr(x, "model_dump"):
                    out.append(x.model_dump())
                elif isinstance(x, dict):
                    out.append(x)
                else:
                    # Try to build dict from iterable; avoid dict(x) when elements aren't (k,v) pairs
                    try:
                        if hasattr(x, "__iter__") and not isinstance(x, str):
                            seq = list(x)
                            if seq and all(isinstance(e, (list, tuple)) and len(e) == 2 for e in seq):
                                out.append(dict(seq))
                            elif seq:
                                # e.g. [(k,v,v), (a,b)]: take first two elements of each
                                out.append({e[0]: e[1] for e in seq if len(e) >= 2})
                            else:
                                out.append({})
                        else:
                            out.append({})
                    except (ValueError, TypeError, IndexError):
                        out.append({})
            return out

        if result is None:
            return {list_key: [], "is_complete": True, "continuation_context": ""}
        if isinstance(result, list):
            return {list_key: ensure_dict_list(result), "is_complete": True, "continuation_context": ""}
        if hasattr(result, "model_dump"):
            d = result.model_dump()
            d[list_key] = ensure_dict_list(d.get(list_key, []))
            return d
        if isinstance(result, dict):
            return {
                list_key: ensure_dict_list(result.get(list_key, [])),
                "is_complete": result.get("is_complete", True),
                "continuation_context": result.get("continuation_context") or ""
            }
        return {list_key: [], "is_complete": True, "continuation_context": ""}

    async def run_summarization(self, novel_text: str, language: str = "Chinese", global_summary: str = None) -> Dict[str, Any]:
        """Step 1: Summarize the novel."""
        context_instruction = ""
        if global_summary:
            context_instruction = f"Global Novel Summary Context:\n{global_summary}\n\nFocus on summarizing this specific part."

        prompt = SUMMARIZER_PROMPT.format(language=language, context_instruction=context_instruction)
        agent = self._create_agent("Summarizer", prompt)
        
        msg = Msg(name="user", content=novel_text, role="user")
        response = await self._call_agent_with_retry(agent, msg, structured_model=NovelSummary)
        result = await self._parse_agent_response(response, NovelSummary)
        # 防御：模型有时会调用 generate_response 但传空参数 {}，导致校验失败；此处对缺失字段做兜底
        if not isinstance(result, dict) or not result.get("title") or not result.get("summary"):
            fallback_title, fallback_summary = self._salvage_novel_summary_from_response(response)
            logger.warning(
                "Summarizer returned invalid/empty tool args; using content fallback. "
                "title=%s, summary_len=%s",
                bool(fallback_title),
                len(fallback_summary or ""),
            )
            result = {
                "title": (result.get("title") if isinstance(result, dict) else None) or fallback_title or "未命名",
                "summary": (result.get("summary") if isinstance(result, dict) else None) or fallback_summary or "",
            }
        return result

    async def run_extraction_fanout(self, 
                                    text: str, 
                                    summary: str, 
                                    language: str = "Chinese",
                                    known_assets: Dict[str, List] = None,
                                    max_iterations: int = 3) -> Dict[str, List]:
        """Step 2: Extract Characters, Scenes, Props in parallel using iterative extraction."""
        known_chars = known_assets.get("characters") if known_assets else None
        known_scenes = known_assets.get("scenes") if known_assets else None
        known_props = known_assets.get("props") if known_assets else None

        # Run extractions in parallel with iteration support
        results = await asyncio.gather(
            self._run_iterative_character_extraction(text, summary, language, known_chars, max_iterations),
            self._run_iterative_scene_extraction(text, summary, language, known_scenes, max_iterations),
            self._run_iterative_prop_extraction(text, summary, language, known_props, max_iterations)
        )
        
        return {
            "characters": results[0],
            "scenes": results[1],
            "props": results[2]
        }
    
    async def _run_iterative_character_extraction(
        self,
        text: str,
        summary: str,
        language: str,
        known_characters: List = None,
        max_iterations: int = 10
    ) -> List[Dict]:
        """Iteratively extract characters with ReAct Agent."""
        all_characters = []
        current_text = text
        iteration = 0
        
        logger.info(f"Starting iterative character extraction for text of length {len(text)}")
        
        while iteration < max_iterations:
            iteration += 1
            logger.info(f"Character extraction iteration {iteration}/{max_iterations}")
            
            # Prepare prompt
            context_char = build_context(summary, known_characters=known_characters)
            gender_zh = ", ".join([f'"{g}"' for g in _OPTIONS.get('character_genders_zh', CHARACTER_GENDERS["zh"])])
            gender_en = ", ".join([f'"{g}"' for g in _OPTIONS.get('character_genders_en', CHARACTER_GENDERS["en"])])
            role_zh = ", ".join([f'"{r}"' for r in _OPTIONS.get('character_roles_zh', CHARACTER_ROLES["zh"])])
            role_en = ", ".join([f'"{r}"' for r in _OPTIONS.get('character_roles_en', CHARACTER_ROLES["en"])])
            emotion_zh = ", ".join([f'"{e}"' for e in _OPTIONS.get('character_emotions_zh', [])])
            emotion_en = ", ".join([f'"{e}"' for e in _OPTIONS.get('character_emotions_en', [])])
            prompt_char = CHARACTER_EXTRACTOR_PROMPT.format(
                language=language,
                context_instruction=context_char,
                gender_zh=gender_zh,
                gender_en=gender_en,
                role_zh=role_zh,
                role_en=role_en,
                emotion_zh=emotion_zh,
                emotion_en=emotion_en
            )
            
            # 创建 agent
            agent = self._create_agent(f"CharacterExtractor_Iter{iteration}", prompt_char)
            
            # Prepare message
            if iteration == 1:
                msg_content = text
            else:
                # Calculate processed position for context
                processed_length = len(text) - len(current_text)
                last_processed_snippet = text[max(0, processed_length - 100):processed_length] if processed_length > 0 else ""
                
                msg_content = f"""Continue extracting characters from the FULL TEXT below.

**IMPORTANT INSTRUCTIONS**:
1. You have already processed the text up to this position: "...{last_processed_snippet}"
2. DO NOT re-extract characters that are already in the list below
3. Continue extraction from where you left off
4. The full text is provided for context to understand character relationships

**Already extracted {len(all_characters)} characters**:
{', '.join([c['name'] for c in all_characters[:20]])}{'...' if len(all_characters) > 20 else ''}

**FULL TEXT** (continue extraction from the marked position):
{text}
"""
            
            msg = Msg(name="user", content=msg_content, role="user")
            
            try:
                response = await self._call_agent_with_retry(agent, msg, structured_model=CharacterList)
                raw_result = await self._parse_agent_response(response, CharacterList)
                result = self._normalize_extraction_result(raw_result, "characters")
                
                iteration_characters = result.get("characters", [])
                is_complete = result.get("is_complete", True)
                continuation_context = result.get("continuation_context", "")
                
                logger.info(f"Iteration {iteration}: Extracted {len(iteration_characters)} characters, is_complete={is_complete}")
                
                if len(iteration_characters) == 0:
                    logger.warning(f"No characters extracted in iteration {iteration}, stopping")
                    break
                
                all_characters.extend(iteration_characters)
                
                if is_complete:
                    logger.info(f"Character extraction completed after {iteration} iterations with {len(all_characters)} total characters")
                    break
                
                if continuation_context:
                    current_text = self._extract_remaining_text(text, continuation_context)
                    if current_text == text or len(current_text.strip()) == 0:
                        break
                else:
                    break
                    
            except Exception as e:
                logger.error(f"Error in character extraction iteration {iteration}: {e}", exc_info=True)
                if iteration == 1:
                    raise
                else:
                    break
        
        return all_characters
    
    async def _run_iterative_scene_extraction(
        self,
        text: str,
        summary: str,
        language: str,
        known_scenes: List = None,
        max_iterations: int = 10
    ) -> List[Dict]:
        """Iteratively extract scenes with ReAct Agent."""
        all_scenes = []
        current_text = text
        iteration = 0
        
        logger.info(f"Starting iterative scene extraction for text of length {len(text)}")
        
        while iteration < max_iterations:
            iteration += 1
            logger.info(f"Scene extraction iteration {iteration}/{max_iterations}")
            
            context_scene = build_context(summary, known_scenes=known_scenes)
            
            env_zh = ", ".join([f'"{e}"' for e in _OPTIONS.get('scene_environments_zh', [])])
            env_en = ", ".join([f'"{e}"' for e in _OPTIONS.get('scene_environments_en', [])])
            time_zh = ", ".join([f'"{t}"' for t in _OPTIONS.get('scene_times_zh', [])])
            time_en = ", ".join([f'"{t}"' for t in _OPTIONS.get('scene_times_en', [])])
            weather_zh = ", ".join([f'"{w}"' for w in _OPTIONS.get('scene_weather_zh', [])])
            weather_en = ", ".join([f'"{w}"' for w in _OPTIONS.get('scene_weather_en', [])])

            prompt_scene = SCENE_EXTRACTOR_PROMPT.format(
                language=language,
                context_instruction=context_scene,
                environment_zh=env_zh,
                environment_en=env_en,
                time_zh=time_zh,
                time_en=time_en,
                weather_zh=weather_zh,
                weather_en=weather_en
            )
            
            agent = self._create_agent(f"SceneExtractor_Iter{iteration}", prompt_scene)
            
            if iteration == 1:
                msg_content = text
            else:
                # Calculate processed position for context
                processed_length = len(text) - len(current_text)
                last_processed_snippet = text[max(0, processed_length - 100):processed_length] if processed_length > 0 else ""
                
                msg_content = f"""Continue extracting scenes from the FULL TEXT below.

**IMPORTANT INSTRUCTIONS**:
1. You have already processed the text up to this position: "...{last_processed_snippet}"
2. DO NOT re-extract scenes that are already in the list below
3. Continue extraction from where you left off
4. The full text is provided for context to understand scene relationships

**Already extracted {len(all_scenes)} scenes**:
{', '.join([s['name'] for s in all_scenes[:20]])}{'...' if len(all_scenes) > 20 else ''}

**FULL TEXT** (continue extraction from the marked position):
{text}
"""
            
            msg = Msg(name="user", content=msg_content, role="user")
            
            try:
                response = await self._call_agent_with_retry(agent, msg, structured_model=SceneList)
                raw_result = await self._parse_agent_response(response, SceneList)
                result = self._normalize_extraction_result(raw_result, "scenes")
                
                iteration_scenes = result.get("scenes", [])
                is_complete = result.get("is_complete", True)
                continuation_context = result.get("continuation_context", "")
                
                logger.info(f"Iteration {iteration}: Extracted {len(iteration_scenes)} scenes, is_complete={is_complete}")
                
                if len(iteration_scenes) == 0:
                    break
                
                all_scenes.extend(iteration_scenes)
                
                if is_complete:
                    logger.info(f"Scene extraction completed after {iteration} iterations with {len(all_scenes)} total scenes")
                    break
                
                if continuation_context:
                    current_text = self._extract_remaining_text(text, continuation_context)
                    if current_text == text or len(current_text.strip()) == 0:
                        break
                else:
                    break
                    
            except Exception as e:
                logger.error(f"Error in scene extraction iteration {iteration}: {e}", exc_info=True)
                if iteration == 1:
                    raise
                else:
                    break
        
        return all_scenes
    
    async def _run_iterative_prop_extraction(
        self,
        text: str,
        summary: str,
        language: str,
        known_props: List = None,
        max_iterations: int = 10
    ) -> List[Dict]:
        """Iteratively extract props with ReAct Agent."""
        all_props = []
        current_text = text
        iteration = 0
        
        logger.info(f"Starting iterative prop extraction for text of length {len(text)}")
        
        while iteration < max_iterations:
            iteration += 1
            logger.info(f"Prop extraction iteration {iteration}/{max_iterations}")
            
            context_prop = build_context(summary, known_props=known_props)
            prompt_prop = PROP_EXTRACTOR_PROMPT.format(
                language=language,
                context_instruction=context_prop
            )
            
            agent = self._create_agent(f"PropExtractor_Iter{iteration}", prompt_prop)
            
            if iteration == 1:
                msg_content = text
            else:
                # Calculate processed position for context
                processed_length = len(text) - len(current_text)
                last_processed_snippet = text[max(0, processed_length - 100):processed_length] if processed_length > 0 else ""
                
                msg_content = f"""Continue extracting props from the FULL TEXT below.

**IMPORTANT INSTRUCTIONS**:
1. You have already processed the text up to this position: "...{last_processed_snippet}"
2. DO NOT re-extract props that are already in the list below
3. Continue extraction from where you left off
4. The full text is provided for context to understand prop usage and significance

**Already extracted {len(all_props)} props**:
{', '.join([p['name'] for p in all_props[:20]])}{'...' if len(all_props) > 20 else ''}

**FULL TEXT** (continue extraction from the marked position):
{text}
"""
            
            msg = Msg(name="user", content=msg_content, role="user")
            
            try:
                response = await self._call_agent_with_retry(agent, msg, structured_model=PropList)
                raw_result = await self._parse_agent_response(response, PropList)
                result = self._normalize_extraction_result(raw_result, "props")
                
                iteration_props = result.get("props", [])
                is_complete = result.get("is_complete", True)
                continuation_context = result.get("continuation_context", "")
                
                logger.info(f"Iteration {iteration}: Extracted {len(iteration_props)} props, is_complete={is_complete}")
                
                if len(iteration_props) == 0:
                    break
                
                all_props.extend(iteration_props)
                
                if is_complete:
                    logger.info(f"Prop extraction completed after {iteration} iterations with {len(all_props)} total props")
                    break
                
                if continuation_context:
                    current_text = self._extract_remaining_text(text, continuation_context)
                    if current_text == text or len(current_text.strip()) == 0:
                        break
                else:
                    break
                    
            except Exception as e:
                logger.error(f"Error in prop extraction iteration {iteration}: {e}", exc_info=True)
                if iteration == 1:
                    raise
                else:
                    break
        
        return all_props

    async def run_deduplication(self, 
                                assets: List[Dict], 
                                asset_type: str, 
                                language: str = "Chinese") -> Dict[str, Any]:
        """Step 3: Deduplicate assets."""
        if not assets or len(assets) < 2:
            return {"assets": assets, "id_mapping": {}}

        items_str = json.dumps(assets, ensure_ascii=False, indent=2)
        prompt = DEDUPLICATOR_PROMPT.format(asset_type=asset_type, language=language)
        agent = self._create_agent(f"Deduplicator_{asset_type}", prompt)
        
        msg = Msg(name="user", content=items_str, role="user")
        response = await self._call_agent_with_retry(agent, msg, structured_model=DeduplicationResult)
        
        res_data = await self._parse_agent_response(response, DeduplicationResult)
        
        final_assets = []
        id_mapping = {}
        
        if isinstance(res_data, dict):
            merged = res_data.get("merged_assets", [])
            mappings = res_data.get("merge_mapping", [])
            final_assets.extend(merged)
            for m in mappings:
                target = m.get("target_id")
                sources = m.get("source_ids", [])
                for s in sources:
                    if s != target:
                        id_mapping[s] = target
        else:
            final_assets = assets

        return {"assets": final_assets, "id_mapping": id_mapping}

    async def run_prompt_optimization(self, 
                                      prompt: str, 
                                      target_type: str = "image",
                                      template: str = "general",
                                      language: str = "Chinese") -> str:
        """Optimize prompt."""
        system_prompt = self._get_optimization_prompt(target_type, template, language)
        agent = self._create_agent("PromptOptimizer", system_prompt)
        
        msg = Msg(name="user", content=prompt, role="user")
        response = await self._call_agent_with_retry(agent, msg)
        
        # 清理 up response content if it's a list (AgentScope change)
        content = response.content
        if isinstance(content, list):
            text_content = ""
            for item in content:
                if isinstance(item, dict) and item.get("type") == "text":
                    text_content += item.get("text", "")
            content = text_content
            
        return content

    async def run_style_recommendation(self,
                                        novel_text: str,
                                        available_styles: List[str]) -> Dict[str, Any]:
        """Recommend an art style based on novel text.

        Args:
            novel_text: The novel text to analyze
            available_styles: List of available style descriptions

        Returns:
            Dictionary with style recommendation
        """
        system_prompt = """You are an expert art director who specializes in visual storytelling.
Your task is to analyze a novel and recommend the most suitable art style for visualizing it.

Based on the novel's genre, tone, themes, and descriptions, recommend ONE style from the available options.
Explain your reasoning briefly.

Available styles:
{styles}

Respond in JSON format:
{{
    "style_name": "Name of the recommended style",
    "reasoning": "Brief explanation of why this style fits the novel"
}}
""".format(styles="\n".join([f"- {s}" for s in available_styles]))

        agent = self._create_agent("StyleRecommender", system_prompt)

        msg = Msg(name="user", content=f"Novel text:\n{novel_text[:5000]}", role="user")
        response = await self._call_agent_with_retry(agent, msg)

        try:
            result = await self._parse_agent_response(response, dict)
            return {
                "style_name": result.get("style_name", ""),
                "reasoning": result.get("reasoning", "")
            }
        except Exception as e:
            logger.error(f"Failed to parse style recommendation: {e}")
            return {
                "style_name": available_styles[0].split(" - ")[0] if available_styles else "General",
                "reasoning": "Default recommendation due to parsing error"
            }

    async def run_script_discussion(self,
                                    novel_text: str,
                                    language: str = "Chinese") -> Dict[str, Any]:
        """Generate adaptation guidance before screenplay drafting."""
        agent = self._create_agent("ScriptDiscussionLead", SCRIPT_DISCUSSION_PROMPT.format(language=language))
        msg = Msg(name="user", content=novel_text[:20000], role="user")
        response = await self._call_agent_with_retry(agent, msg, structured_model=ScriptDiscussion)
        result = await self._parse_agent_response(response, ScriptDiscussion)
        if isinstance(result, dict):
            return result
        return {
            "adaptation_strategy": "",
            "episode_structure": [],
            "character_focus": [],
            "dialogue_guidelines": [],
        }

    async def run_screenplay_adaptation(self,
                                        source_text: str,
                                        global_summary: str,
                                        adaptation_strategy: str,
                                        dialogue_guidelines: List[str],
                                        language: str = "Chinese") -> Dict[str, Any]:
        """Adapt prose into an editable screenplay draft."""
        system_prompt = SCREENPLAY_ADAPTATION_PROMPT.format(
            global_summary=global_summary or "无",
            adaptation_strategy=adaptation_strategy or "忠于原著，优先可编辑性。",
            dialogue_guidelines="\n".join(f"- {item}" for item in (dialogue_guidelines or [])) or "- 保持对白自然。",
            language=language,
        )
        agent = self._create_agent("ScreenplayWriter", system_prompt)
        msg = Msg(name="user", content=source_text[:20000], role="user")
        response = await self._call_agent_with_retry(agent, msg, structured_model=ScreenplayDraft)
        result = await self._parse_agent_response(response, ScreenplayDraft)
        if isinstance(result, dict):
            return result
        return {
            "title": "未命名剧本",
            "summary": "",
            "screenplay": source_text,
        }

    def _extract_remaining_text(self, original_text: str, continuation_context: str) -> str:
        """ 提取 the remaining text to process based on continuation context.
        
        Args:
            original_text: The original full text
            continuation_context: Context string from the previous iteration
            
        Returns:
            The remaining text to process
        """
        try:
            # 1. Try to find "Next text to process" marker (English or Chinese)
            next_markers = ["Next text to process:", "下一步要处理的文本：", "下一步要处理的文本:"]
            for marker in next_markers:
                if marker in continuation_context:
                    # 提取 the text after this marker
                    parts = continuation_context.split(marker)
                    if len(parts) > 1:
                        # 获取 the line immediately following the marker
                        rest_of_string = parts[1].strip()
                        first_line = rest_of_string.split("\n")[0].strip()
                        
                        # 清理 up quotes from the line
                        next_text_marker = first_line.strip('"').strip("'")
                        
                        # 查找 this text in the original
                        # 注意: We should search in the original text, but we need to be careful about duplicates.
                        # Ideally we should search near where we think we are, but simple search is a start.
                        if next_text_marker and next_text_marker in original_text:
                            idx = original_text.index(next_text_marker)
                            remaining = original_text[idx:]
                            logger.info(f"Extracted remaining text starting from: {next_text_marker[:50]}...")
                            return remaining
            
            # 2. Fallback: Try to find "Last processed text" marker
            last_markers = ["Last processed text:", "最后处理的文本：", "最后处理的文本:"]
            for marker in last_markers:
                if marker in continuation_context:
                    parts = continuation_context.split(marker)
                    if len(parts) > 1:
                        # 获取 the line immediately following the marker
                        rest_of_string = parts[1].strip()
                        first_line = rest_of_string.split("\n")[0].strip()
                        
                        # 清理 up quotes from the line
                        last_text_marker = first_line.strip('"').strip("'")
                        
                        # 查找 this text in the original and get everything after it
                        if last_text_marker and last_text_marker in original_text:
                            idx = original_text.index(last_text_marker)
                            # Move past the last processed text
                            idx += len(last_text_marker)
                            remaining = original_text[idx:].strip()
                            logger.info(f"Extracted remaining text after: {last_text_marker[:50]}...")
                            return remaining
            
            # 3. If we can't parse the continuation context, return the original text
            logger.warning(f"Could not parse continuation_context: {continuation_context[:100]}..., returning original text")
            return original_text
            
        except Exception as e:
            logger.error(f"Error extracting remaining text: {e}", exc_info=True)
            return original_text

    async def run_storyboard_splitting(self,
                                       novel_text: str,
                                       language: str = "Chinese",
                                       known_assets: Dict[str, List] = None,
                                       max_iterations: int = 10) -> Dict[str, Any]:
        """Split storyboards with integrated prompts using iterative ReAct Agent.
        
        Args:
            novel_text: The novel text to split into storyboards
            language: Target language
            known_assets: Dictionary of known characters, scenes, props
            max_iterations: Maximum number of iterations to prevent infinite loops
            
        Returns:
            Dictionary containing all storyboards
        """
        known_chars = known_assets.get("characters") if known_assets else None
        known_scenes = known_assets.get("scenes") if known_assets else None
        known_props = known_assets.get("props") if known_assets else None
        
        # 构建 context string with image_prompt information
        context_parts = []
        
        if known_chars:
            context_parts.append("**Characters:**")
            for char in known_chars:
                char_info = f"- {char.get('name', 'Unknown')}: {char.get('desc', '')}"
                if char.get('appearance'):
                    char_info += f"\n  Appearance: {char['appearance']}"
                if char.get('image_prompt'):
                    char_info += f"\n  Visual Prompt: {char['image_prompt']}"
                context_parts.append(char_info)
        
        if known_scenes:
            context_parts.append("\n**Scenes:**")
            for scene in known_scenes:
                scene_info = f"- {scene.get('name', 'Unknown')}: {scene.get('desc', '')}"
                if scene.get('location'):
                    scene_info += f"\n  Location: {scene['location']}"
                if scene.get('atmosphere'):
                    scene_info += f"\n  Atmosphere: {scene['atmosphere']}"
                if scene.get('image_prompt'):
                    scene_info += f"\n  Visual Prompt: {scene['image_prompt']}"
                context_parts.append(scene_info)
        
        if known_props:
            context_parts.append("\n**Props:**")
            for prop in known_props:
                prop_info = f"- {prop.get('name', 'Unknown')}: {prop.get('desc', '')}"
                if prop.get('image_prompt'):
                    prop_info += f"\n  Visual Prompt: {prop['image_prompt']}"
                context_parts.append(prop_info)
        
        context_str = "\n".join(context_parts) if context_parts else "No existing assets provided."
        
        lang_key = "zh" if language in ["Chinese", "zh"] else "en"
        shot_types_list = _OPTIONS.get(f'shot_types_{lang_key}', SHOT_TYPES.get(lang_key, SHOT_TYPES["en"]))
        shot_types_str = ", ".join([f'"{s}"' for s in shot_types_list])
        camera_movements_list = _OPTIONS.get(f'camera_movements_{lang_key}', CAMERA_MOVEMENTS.get(lang_key, CAMERA_MOVEMENTS["en"]))
        camera_movements_str = ", ".join([f'"{c}"' for c in camera_movements_list])
        transitions_list = _OPTIONS.get(f'transitions_{lang_key}', TRANSITIONS.get(lang_key, TRANSITIONS["en"]))
        transitions_str = ", ".join([f'"{t}"' for t in transitions_list])
        
        # New cinematic options
        camera_angles_list = _OPTIONS.get(f'camera_angles_{lang_key}', [])
        camera_angles_str = ", ".join([f'"{a}"' for a in camera_angles_list])
        lenses_list = _OPTIONS.get(f'lenses_{lang_key}', [])
        lenses_str = ", ".join([f'"{l}"' for l in lenses_list])
        focus_list = _OPTIONS.get(f'focus_{lang_key}', [])
        focus_str = ", ".join([f'"{f}"' for f in focus_list])
        lighting_list = _OPTIONS.get(f'lighting_{lang_key}', [])
        lighting_str = ", ".join([f'"{l}"' for l in lighting_list])
        color_style_list = _OPTIONS.get(f'color_style_{lang_key}', [])
        color_style_str = ", ".join([f'"{c}"' for c in color_style_list])

        system_prompt = STORYBOARD_ARTIST_PROMPT.format(
            context_str=context_str,
            language=language,
            shot_types_str=shot_types_str,
            camera_movements_str=camera_movements_str,
            camera_angles_str=camera_angles_str,
            lenses_str=lenses_str,
            focus_str=focus_str,
            lighting_str=lighting_str,
            color_style_str=color_style_str,
            transitions_str=transitions_str
        )
        
        # 初始化 collection for all storyboards
        all_storyboards = []
        current_text = novel_text
        iteration = 0
        
        logger.info(f"Starting iterative storyboard generation for text of length {len(novel_text)}")
        
        while iteration < max_iterations:
            iteration += 1
            logger.info(f"Storyboard generation iteration {iteration}/{max_iterations}, processing text length: {len(current_text)}")
            
            # 创建 agent for this iteration
            agent = self._create_agent(f"StoryboardArtist_Iter{iteration}", system_prompt)
            
            # Prepare message with continuation context if this is not the first iteration
            if iteration == 1:
                msg_content = current_text
            else:
                msg_content = f"""Continue generating storyboards from where you left off.

**Previous Progress:**
- Already generated: {len(all_storyboards)} shots
- Last shot number: {all_storyboards[-1]['shot_number']}
- Last shot title: {all_storyboards[-1].get('shot_title', 'N/A')}

**Instructions:**
1. Start shot numbering from {len(all_storyboards) + 1}
2. Continue the narrative flow naturally
3. Process as much of the remaining text as possible
4. Set is_complete=true ONLY if you finish ALL remaining content
5. Set is_complete=false if you need another iteration, and provide clear continuation_context

**Remaining Text to Process:**
{current_text}
"""
            
            msg = Msg(name="user", content=msg_content, role="user")
            
            try:
                response = await self._call_agent_with_retry(agent, msg, structured_model=StoryboardList)
                result = await self._parse_agent_response(response, StoryboardList)
                
                # 提取 storyboards from this iteration
                iteration_storyboards = result.get("storyboards", [])
                is_complete = result.get("is_complete", True)
                continuation_context = result.get("continuation_context", "")
                
                logger.info(f"Iteration {iteration}: Generated {len(iteration_storyboards)} shots, is_complete={is_complete}")
                
                if len(iteration_storyboards) == 0:
                    logger.warning(f"No storyboards generated in iteration {iteration}, stopping")
                    break
                
                # Add to collection
                all_storyboards.extend(iteration_storyboards)
                
                # Check if complete
                if is_complete:
                    logger.info(f"Storyboard generation completed after {iteration} iterations with {len(all_storyboards)} total shots")
                    break
                
                # Not complete, extract the remaining text to process
                if continuation_context:
                    logger.info(f"Continuation context received: {continuation_context[:300]}...")
                    
                    # 提取 remaining text based on continuation context
                    remaining_text = self._extract_remaining_text(novel_text, continuation_context)
                    
                    if remaining_text == current_text:
                        # Couldn't extract new text, we might be stuck in a loop
                        logger.warning("Could not extract new remaining text, stopping to prevent infinite loop")
                        break
                    
                    current_text = remaining_text
                    
                    if len(current_text.strip()) == 0:
                        logger.info("No remaining text to process, marking as complete")
                        break
                    
                    if iteration >= max_iterations:
                        logger.warning(f"Reached max iterations ({max_iterations}), stopping generation")
                        break
                else:
                    logger.warning("is_complete=False but no continuation_context provided, stopping")
                    break
                    
            except Exception as e:
                logger.error(f"Error in storyboard generation iteration {iteration}: {e}", exc_info=True)
                if iteration == 1:
                    # First iteration fails, raise the error
                    raise
                else:
                    # Later iterations fail, return what we have so far
                    logger.warning(f"Stopping at iteration {iteration} due to error, returning {len(all_storyboards)} shots")
                    break
        
        return {
            "storyboards": all_storyboards,
            "is_complete": iteration < max_iterations and is_complete,
            "total_iterations": iteration
        }

    async def run_chapter_splitting(self, novel_text: str, language: str = "Chinese") -> Dict[str, Any]:
        """Split novel text into chapters using agent."""
        prompt = CHAPTER_SPLITTER_PROMPT.format(language=language, context_instruction="")
        agent = self._create_agent("ChapterSplitter", prompt)
        
        # Truncate if too long (simple safety, ideally should chunk)
        # 假设这是用于较短文本或模型处理长上下文的情况
        msg = Msg(name="user", content=novel_text, role="user")
        response = await self._call_agent_with_retry(agent, msg, structured_model=ChapterList)
        return await self._parse_agent_response(response, ChapterList)

    def _get_optimization_prompt(self, target_type, template, language):
        """ 获取 optimization system prompt based on target type and template."""
        if target_type == "image":
            if template == "storyboard":
                system_prompt = OPTIMIZATION_PROMPT_IMAGE_STORYBOARD
            elif template == "character_portrait":
                system_prompt = OPTIMIZATION_PROMPT_IMAGE_WHITE_BG
            elif template == "character_three_view":
                system_prompt = OPTIMIZATION_PROMPT_IMAGE_THREE_VIEW
            else:
                system_prompt = OPTIMIZATION_PROMPT_IMAGE_DEFAULT
        elif target_type == "video":
            if template == "storyboard":
                system_prompt = OPTIMIZATION_PROMPT_VIDEO_STORYBOARD_INTEGRATED
            elif template == "asset_360":
                system_prompt = OPTIMIZATION_PROMPT_VIDEO_ASSET_360
            else:
                system_prompt = OPTIMIZATION_PROMPT_VIDEO_DEFAULT
        else:
            system_prompt = OPTIMIZATION_PROMPT_DEFAULT
            
        return system_prompt.format(language=language)
