"""LangGraph workflow for Opus Orchestrator - FIXED VERSION.

Key fixes based on Gemini's analysis:
1. Nodes return dicts instead of mutating state
2. run() uses stream_mode="values" 
3. Falls back to get_state() from checkpointer
"""

import os
from typing import Any, Optional

from dotenv import load_dotenv

load_dotenv("/home/solaria/.openclaw/workspace/opus-orchestrator-ai/.env")

from pydantic import BaseModel, Field
from enum import Enum

from langgraph.graph import StateGraph, END
from langgraph.checkpoint.memory import MemorySaver

from opus_orchestrator.frameworks import get_framework_prompt, StoryFramework
from opus_orchestrator.utils.llm_sync import LLMClient


# ============== STATE SCHEMA ==============

class Stage(str, Enum):
    """Workflow stages."""
    SEED = "seed"
    ONE_SENTENCE = "one_sentence"
    ONE_PARAGRAPH = "one_paragraph"
    CHARACTER_SHEETS = "character_sheets"
    FOUR_PAGE_OUTLINE = "four_page_outline"
    CHARACTER_CHARTS = "character_charts"
    SCENE_LIST = "scene_list"
    SCENE_DESCRIPTIONS = "scene_descriptions"
    STYLE_GUIDE = "style_guide"
    WRITING = "writing"
    COMPLETE = "complete"


class Character(BaseModel):
    """Character schema."""
    name: str = ""
    role: str = ""
    description: str = ""
    want: str = ""
    need: str = ""
    fear: str = ""
    arc: str = ""


class PlotBeat(BaseModel):
    """Scene/beat schema."""
    name: str = ""
    description: str = ""


class ChapterPlan(BaseModel):
    """Chapter plan schema."""
    chapter_number: int = 0
    title: str = ""
    summary: str = ""
    word_count_target: int = 3000


class PreWriting(BaseModel):
    """Pre-writing output schema."""
    one_sentence: str = ""
    one_paragraph: str = ""
    characters: list[Character] = Field(default_factory=list)
    outline_sections: list[str] = Field(default_factory=list)
    character_details: dict[str, str] = Field(default_factory=dict)
    scene_list: list[PlotBeat] = Field(default_factory=list)
    chapter_plans: list[ChapterPlan] = Field(default_factory=list)
    scene_descriptions: dict[str, str] = Field(default_factory=dict)
    framework_used: str = "snowflake"


class ChapterState(BaseModel):
    """Chapter writing state."""
    content: str = ""
    word_count: int = 0
    critique_score: float = 0.0
    iterations: int = 0
    approved: bool = False


class OpusGraphState(BaseModel):
    """Main state for LangGraph."""
    stage: Stage = Stage.SEED
    framework: str = "snowflake"
    genre: str = "general"
    target_word_count: int = 80000
    seed_concept: str = ""
    
    prewriting: PreWriting = Field(default_factory=PreWriting)
    style_guide: str = ""
    
    current_chapter: int = 0
    chapters: dict[int, ChapterState] = Field(default_factory=dict)
    
    manuscript: str = ""
    total_word_count: int = 0
    
    validation_errors: list[str] = Field(default_factory=list)
    warnings: list[str] = Field(default_factory=list)
    
    progress: float = 0.0
    messages: list[str] = Field(default_factory=list)


# ============== WORKFLOW ==============

class OpusGraph:
    """LangGraph workflow - FIXED with dict returns."""
    
    def __init__(
        self,
        framework: str = "snowflake",
        genre: str = "general",
        target_word_count: int = 80000,
        api_key: Optional[str] = None,
    ):
        self.framework = framework
        self.genre = genre
        self.target_word_count = target_word_count
        self.api_key = api_key or os.environ.get("OPENAI_API_KEY")
        
        # Use synchronous LLM
        self.llm = LLMClient(api_key=self.api_key, provider="openai", model="gpt-4o")
        
        # Build graph
        self.graph = self._build_graph()
    
    def _call_llm(self, system_prompt: str, user_prompt: str) -> str:
        """Call LLM synchronously."""
        return self.llm.complete(system_prompt, user_prompt)
    
    def _build_graph(self) -> StateGraph:
        """Build the LangGraph."""
        
        workflow = StateGraph(OpusGraphState)
        
        # Add nodes - each returns a dict
        workflow.add_node("seed", self.node_seed)
        workflow.add_node("one_sentence", self.node_one_sentence)
        workflow.add_node("one_paragraph", self.node_one_paragraph)
        workflow.add_node("character_sheets", self.node_character_sheets)
        workflow.add_node("four_page_outline", self.node_four_page_outline)
        workflow.add_node("character_charts", self.node_character_charts)
        workflow.add_node("scene_list", self.node_scene_list)
        workflow.add_node("scene_descriptions", self.node_scene_descriptions)
        workflow.add_node("style_guide", self.node_style_guide)
        workflow.add_node("write_chapters", self.node_write_chapters)
        workflow.add_node("complete", self.node_complete)
        
        # Edges
        workflow.set_entry_point("seed")
        workflow.add_edge("seed", "one_sentence")
        workflow.add_edge("one_sentence", "one_paragraph")
        workflow.add_edge("one_paragraph", "character_sheets")
        workflow.add_edge("character_sheets", "four_page_outline")
        workflow.add_edge("four_page_outline", "character_charts")
        workflow.add_edge("character_charts", "scene_list")
        workflow.add_edge("scene_list", "scene_descriptions")
        workflow.add_edge("scene_descriptions", "style_guide")
        workflow.add_edge("style_guide", "write_chapters")
        workflow.add_edge("write_chapters", "complete")
        workflow.add_edge("complete", END)
        
        checkpointer = MemorySaver()
        return workflow.compile(checkpointer=checkpointer)
    
    # ============== NODES (Return DICT, not mutated state) ==============
    
    def node_seed(self, state: OpusGraphState) -> dict:
        """Initialize from seed."""
        print(f"\n🌱 SEED: {state.seed_concept[:80]}...")
        
        return {
            "stage": Stage.ONE_SENTENCE,
            "progress": 0.05,
            "messages": [f"Started: {state.seed_concept[:50]}"],
        }
    
    def node_one_sentence(self, state: OpusGraphState) -> dict:
        """Stage 1: One sentence."""
        print("📝 STAGE 1: One sentence...")
        
        system_prompt = get_framework_prompt(StoryFramework(self.framework))
        user_prompt = f"""Create ONE SENTENCE that captures this story.

Must include: Protagonist, Goal, Conflict, Stakes

Seed: {state.seed_concept}
"""
        
        result = self._call_llm(system_prompt, user_prompt)
        
        # Update prewriting via dict return
        new_prewriting = state.prewriting.model_copy()
        new_prewriting.one_sentence = result.strip()
        
        return {
            "prewriting": new_prewriting,
            "stage": Stage.ONE_SENTENCE,
            "progress": 0.12,
            "messages": state.messages + [f"One sentence: {result.strip()[:60]}..."],
        }
    
    def node_one_paragraph(self, state: OpusGraphState) -> dict:
        """Stage 2: One paragraph."""
        print("📝 STAGE 2: One paragraph...")
        
        system_prompt = get_framework_prompt(StoryFramework(self.framework))
        user_prompt = f"""Expand to ONE PARAGRAPH (4-8 sentences):

Include: Opening, Setup, Catalyst, Rising Action, Midpoint, Complications, Crisis, Resolution

One sentence: {state.prewriting.one_sentence}
"""
        
        result = self._call_llm(system_prompt, user_prompt)
        
        new_prewriting = state.prewriting.model_copy()
        new_prewriting.one_paragraph = result.strip()
        
        return {
            "prewriting": new_prewriting,
            "stage": Stage.ONE_PARAGRAPH,
            "progress": 0.20,
            "messages": state.messages + ["One paragraph complete"],
        }
    
    def node_character_sheets(self, state: OpusGraphState) -> dict:
        """Stage 3: Character sheets."""
        print("📝 STAGE 3: Character sheets...")
        
        system_prompt = "You are a character development expert."
        user_prompt = f"""Create character sheets.

For each: Name, Role, Want, Need, Fear

Story: {state.prewriting.one_paragraph}
"""
        
        result = self._call_llm(system_prompt, user_prompt)
        
        characters = self._parse_characters(result)
        
        new_prewriting = state.prewriting.model_copy()
        new_prewriting.characters = characters
        
        return {
            "prewriting": new_prewriting,
            "stage": Stage.CHARACTER_SHEETS,
            "progress": 0.30,
            "messages": state.messages + [f"Created {len(characters)} characters"],
        }
    
    def node_four_page_outline(self, state: OpusGraphState) -> dict:
        """Stage 4: Four-page outline."""
        print("📝 STAGE 4: Four-page outline...")
        
        system_prompt = get_framework_prompt(StoryFramework(self.framework))
        user_prompt = f"""Create a detailed outline.

Story: {state.prewriting.one_paragraph}
Characters: {', '.join(c.name for c in state.prewriting.characters)}
"""
        
        result = self._call_llm(system_prompt, user_prompt)
        
        new_prewriting = state.prewriting.model_copy()
        new_prewriting.outline_sections = [s.strip() for s in result.split("\n\n") if s.strip()]
        
        return {
            "prewriting": new_prewriting,
            "stage": Stage.FOUR_PAGE_OUTLINE,
            "progress": 0.40,
            "messages": state.messages + ["Outline complete"],
        }
    
    def node_character_charts(self, state: OpusGraphState) -> dict:
        """Stage 5: Character charts."""
        print("📝 STAGE 5: Character charts...")
        
        system_prompt = "You are a character development expert."
        user_prompt = f"""Create detailed character profiles.

Characters: {', '.join(c.name for c in state.prewriting.characters)}

Include: Backstory, Psychology, Speech patterns, Key scenes
"""
        
        result = self._call_llm(system_prompt, user_prompt)
        
        new_prewriting = state.prewriting.model_copy()
        for char in new_prewriting.characters:
            new_prewriting.character_details[char.name] = result[:800]
        
        return {
            "prewriting": new_prewriting,
            "stage": Stage.CHARACTER_CHARTS,
            "progress": 0.50,
            "messages": state.messages + ["Character charts complete"],
        }
    
    def node_scene_list(self, state: OpusGraphState) -> dict:
        """Stage 6: Scene list."""
        print("📝 STAGE 6: Scene list...")
        
        num_scenes = max(10, self.target_word_count // 1500)
        
        system_prompt = get_framework_prompt(StoryFramework(self.framework))
        user_prompt = f"""Create {num_scenes} scenes.

For each: name, description, POV, location
"""
        
        result = self._call_llm(system_prompt, user_prompt)
        
        scenes = self._parse_scenes(result)
        
        # Create chapter plans
        num_chapters = max(3, self.target_word_count // 3000)
        scenes_per_ch = max(1, len(scenes) // num_chapters)
        
        chapter_plans = []
        for i in range(num_chapters):
            start = i * scenes_per_ch
            end = min(start + scenes_per_ch, len(scenes))
            chapter_plans.append(ChapterPlan(
                chapter_number=i + 1,
                title=f"Chapter {i + 1}",
                summary=f"Chapter {i + 1}",
                word_count_target=self.target_word_count // num_chapters,
            ))
        
        new_prewriting = state.prewriting.model_copy()
        new_prewriting.scene_list = scenes
        new_prewriting.chapter_plans = chapter_plans
        
        return {
            "prewriting": new_prewriting,
            "stage": Stage.SCENE_LIST,
            "progress": 0.60,
            "messages": state.messages + [f"{len(scenes)} scenes, {num_chapters} chapters"],
        }
    
    def node_scene_descriptions(self, state: OpusGraphState) -> dict:
        """Stage 7: Scene descriptions."""
        print("📝 STAGE 7: Scene descriptions...")
        
        system_prompt = "You are a story architect."
        user_prompt = "Describe key scenes."
        
        result = self._call_llm(system_prompt, user_prompt)
        
        new_prewriting = state.prewriting.model_copy()
        new_prewriting.scene_descriptions = {"key_scenes": result[:2000]}
        
        return {
            "prewriting": new_prewriting,
            "stage": Stage.SCENE_DESCRIPTIONS,
            "progress": 0.70,
            "messages": state.messages + ["Scene descriptions complete"],
        }
    
    def node_style_guide(self, state: OpusGraphState) -> dict:
        """Create style guide."""
        print("🎨 STYLE GUIDE...")
        
        system_prompt = "You are a prose style expert."
        user_prompt = f"""Create a style guide.

Genre: {self.genre}
"""
        
        result = self._call_llm(system_prompt, user_prompt)
        
        return {
            "style_guide": result.strip(),
            "stage": Stage.STYLE_GUIDE,
            "progress": 0.75,
            "messages": state.messages + ["Style guide created"],
        }
    
    def node_write_chapters(self, state: OpusGraphState) -> dict:
        """Write all chapters."""
        print("\n✍️  WRITING CHAPTERS...")
        
        system_prompt = f"""You are a professional novelist.
Style: {state.style_guide[:500] if state.style_guide else 'Professional fiction'}
"""
        
        chapters = {}
        
        for plan in state.prewriting.chapter_plans:
            print(f"   Writing chapter {plan.chapter_number}...")
            
            user_prompt = f"""Write Chapter {plan.chapter_number}: {plan.summary}

Story: {state.prewriting.one_sentence}
Characters: {', '.join(c.name for c in state.prewriting.characters[:3])}

Write ~{plan.word_count_target} words.
"""
            
            result = self._call_llm(system_prompt, user_prompt)
            
            chapters[plan.chapter_number] = ChapterState(
                content=result.strip(),
                word_count=len(result.split()),
                critique_score=0.8,
                iterations=1,
                approved=True,
            )
            
            print(f"      → {len(result.split())} words")
        
        return {
            "chapters": chapters,
            "stage": Stage.WRITING,
            "progress": 0.90,
            "messages": state.messages + [f"Wrote {len(chapters)} chapters"],
        }
    
    def node_complete(self, state: OpusGraphState) -> dict:
        """Complete - compile manuscript."""
        # Compile manuscript
        parts = []
        for i in range(1, len(state.chapters) + 1):
            if i in state.chapters:
                parts.append(f"# Chapter {i}\n\n{state.chapters[i].content}")
        
        manuscript = "\n\n---\n\n".join(parts)
        total_words = sum(c.word_count for c in state.chapters.values())
        
        print(f"\n✅ COMPLETE!")
        print(f"   Chapters: {len(state.chapters)}")
        print(f"   Words: {total_words:,}")
        
        return {
            "manuscript": manuscript,
            "total_word_count": total_words,
            "stage": Stage.COMPLETE,
            "progress": 1.0,
            "messages": state.messages + [f"Final: {total_words} words"],
        }
    
    # ============== PARSING ==============
    
    def _parse_characters(self, text: str) -> list[Character]:
        """Parse characters from text."""
        characters = []
        
        for line in text.split("\n"):
            line = line.strip()
            lower = line.lower()
            
            if "name:" in lower and len(line) < 50:
                name = line.split(":", 1)[-1].strip()
                characters.append(Character(
                    name=name,
                    role="character",
                    description=line,
                    want="To be defined",
                    need="To be defined",
                    fear="Unknown",
                ))
        
        if not characters:
            characters.append(Character(
                name="Protagonist",
                role="protagonist",
                description="Main character",
                want="Complete the quest",
                need="Learn and grow",
                fear="Failure",
            ))
        
        return characters[:5]
    
    def _parse_scenes(self, text: str) -> list[PlotBeat]:
        """Parse scenes from text."""
        scenes = []
        
        for i, line in enumerate(text.split("\n")):
            line = line.strip()
            if line and len(line) > 10:
                scenes.append(PlotBeat(
                    name=f"Scene {i+1}",
                    description=line[:120],
                ))
        
        return scenes[:20] if scenes else [PlotBeat(name=f"Scene {i+1}", description=f"Beat {i+1}") for i in range(10)]
    
    # ============== RUN (FIXED) ==============
    
    def run(self, seed_concept: str, thread_id: str = "default") -> OpusGraphState:
        """Run the workflow - FIXED with stream_mode='values'."""
        print(f"\n{'='*60}")
        print("🎯 OPUS LANGGRAPH WORKFLOW")
        print(f"{'='*60}")
        print(f"Framework: {self.framework}")
        print(f"Target: {self.target_word_count:,} words\n")
        
        initial_state = OpusGraphState(
            seed_concept=seed_concept,
            framework=self.framework,
            genre=self.genre,
            target_word_count=self.target_word_count,
        )
        
        config = {"configurable": {"thread_id": thread_id}}
        
        # FIX: Use stream_mode="values" to get full state objects
        final_state = initial_state
        
        try:
            # Stream with values mode - each chunk IS the full state
            for chunk in self.graph.stream(initial_state, config, stream_mode="values"):
                if isinstance(chunk, OpusGraphState):
                    final_state = chunk
                    # Print progress
                    if chunk.messages:
                        last_msg = chunk.messages[-1]
                        if "Writing chapter" in last_msg or "COMPLETE" in last_msg:
                            print(last_msg)
        except Exception as e:
            print(f"Stream error: {e}")
        
        # FIX: Fallback - pull directly from checkpointer
        if not final_state.manuscript:
            print("Pulling from checkpointer...")
            try:
                snapshot = self.graph.get_state(config)
                if snapshot and snapshot.values:
                    final_state = snapshot.values
            except Exception as e:
                print(f"Checkpointer error: {e}")
        
        print(f"\n[RESULT] Chapters: {len(final_state.chapters)}, Words: {final_state.total_word_count}")
        
        return final_state


def run_opus(
    seed_concept: str,
    framework: str = "snowflake",
    genre: str = "general",
    target_word_count: int = 80000,
    thread_id: str = "default",
) -> OpusGraphState:
    """Convenience function."""
    api_key = os.environ.get("OPENAI_API_KEY")
    
    workflow = OpusGraph(
        framework=framework,
        genre=genre,
        target_word_count=target_word_count,
        api_key=api_key,
    )
    
    return workflow.run(seed_concept, thread_id)