import streamlit as st
from langchain import PromptTemplate
from typing import TypedDict, List, Dict, Optional
from langchain.graphs import StateGraph
from dataclasses import dataclass, field
import random

# Data Structures
@dataclass
class StoryState:
    current_step: int = 0
    max_steps: int = 5
    story_log: List[str] = field(default_factory=list)
    user_inputs: List[str] = field(default_factory=list)
    character1_responses: List[str] = field(default_factory=list)
    character2_responses: List[str] = field(default_factory=list)
    story_outcome: Optional[str] = None

class MicroStory(TypedDict):
    title: str
    initial_setup: str
    character1_name: str
    character2_name: str
    steps: List[str]
    success_conditions: List[str]
    failure_conditions: List[str]

from langchain.graphs import StateGraph
from typing import Dict, List, Any
from dataclasses import dataclass
from enum import Enum

class StoryNodeType(Enum):
    SETUP = "setup"
    USER_INPUT = "user_input"
    CHARACTER1_RESPONSE = "character1_response"
    CHARACTER2_RESPONSE = "character2_response"
    EVALUATION = "evaluation"

@dataclass
class StoryGraphState:
    current_node: StoryNodeType
    story_data: Dict[str, Any]
    accumulated_context: List[Dict[str, str]]
    step_count: int = 0

def create_story_graph() -> StateGraph:
    """Creates the state graph for story progression"""
    graph = StateGraph()

    # Define state transitions
    def setup_to_user_input(state: StoryGraphState) -> StoryGraphState:
        state.current_node = StoryNodeType.USER_INPUT
        return state

    def user_input_to_char1(state: StoryGraphState, user_input: str) -> StoryGraphState:
        state.current_node = StoryNodeType.CHARACTER1_RESPONSE
        state.accumulated_context.append({"role": "user", "content": user_input})
        return state

    def char1_to_char2(state: StoryGraphState, char1_response: str) -> StoryGraphState:
        state.current_node = StoryNodeType.CHARACTER2_RESPONSE
        state.accumulated_context.append({"role": "character1", "content": char1_response})
        return state

    def char2_to_evaluation(state: StoryGraphState, char2_response: str) -> StoryGraphState:
        state.current_node = StoryNodeType.EVALUATION
        state.accumulated_context.append({"role": "character2", "content": char2_response})
        state.step_count += 1
        return state

    def evaluation_to_next(state: StoryGraphState) -> StoryGraphState:
        if state.step_count >= 5:
            # Story is complete, stay in evaluation
            return state
        # Move to next user input
        state.current_node = StoryNodeType.USER_INPUT
        return state

    # Add nodes and edges
    graph.add_node("setup", setup_to_user_input)
    graph.add_node("user_input", user_input_to_char1)
    graph.add_node("character1_response", char1_to_char2)
    graph.add_node("character2_response", char2_to_evaluation)
    graph.add_node("evaluation", evaluation_to_next)

    # Connect nodes
    graph.add_edge("setup", "user_input")
    graph.add_edge("user_input", "character1_response")
    graph.add_edge("character1_response", "character2_response")
    graph.add_edge("character2_response", "evaluation")
    graph.add_edge("evaluation", "user_input")

    return graph

class StoryRunner:
    def __init__(self, story_data: Dict[str, Any]):
        self.graph = create_story_graph()
        self.state = StoryGraphState(
            current_node=StoryNodeType.SETUP,
            story_data=story_data,
            accumulated_context=[]
        )

    def process_user_input(self, user_input: str) -> Dict[str, Any]:
        """Process user input and advance the story state"""
        if self.state.current_node != StoryNodeType.USER_INPUT:
            raise ValueError("Not ready for user input")

        # Advance through the graph
        self.state = self.graph.transition("user_input", self.state, user_input)
        self.state = self.graph.transition("character1_response", self.state, None)
        self.state = self.graph.transition("character2_response", self.state, None)
        self.state = self.graph.transition("evaluation", self.state, None)

        # Return current state info
        return {
            "step_count": self.state.step_count,
            "is_complete": self.state.step_count >= 5,
            "current_context": self.state.accumulated_context[-3:] if self.state.accumulated_context else [],
            "current_node": self.state.current_node.value
        }

    def get_full_context(self) -> List[Dict[str, str]]:
        """Get the full conversation context"""
        return self.state.accumulated_context

    def is_complete(self) -> bool:
        """Check if the story is complete"""
        return self.state.step_count >= 5

# Sample Stories Database
# Story Categories and Templates
STORY_CATEGORIES = {
    "Mystery": [
        {
            "title": "The Library Mystery",
            "initial_setup": "In the ancient library of St. Bartholomew's, a rare manuscript has gone missing.",
            "character1_name": "Detective Nash",
            "character2_name": "Librarian Wells",
            "steps": [
                "You notice strange symbols carved into the reading desk",
                "A student mentions seeing someone in medieval clothing",
                "The manuscript tracking system shows impossible timestamps",
                "Temperature drops significantly in the rare books section",
                "You find a hidden door behind the card catalog"
            ],
            "success_conditions": [
                "mentioned checking the security cameras",
                "investigated the symbols",
                "questioned the student further",
                "connected medieval sighting with timestamps"
            ],
            "failure_conditions": [
                "accused the librarian",
                "ignored the symbols",
                "left the library",
                "called the police immediately"
            ]
        },
        {
            "title": "The Digital Deception",
            "initial_setup": "A tech startup's revolutionary AI algorithm has been stolen right before a major demo.",
            "character1_name": "Cyber Detective Chen",
            "character2_name": "System Admin Rodriguez",
            "steps": [
                "The server logs show multiple failed login attempts",
                "An employee reports receiving a suspicious email",
                "The backup system was manually disabled",
                "Strange network traffic appears during off-hours",
                "A hidden backdoor program is discovered"
            ],
            "success_conditions": [
                "checked email headers",
                "analyzed network logs",
                "investigated backup system",
                "traced the backdoor"
            ],
            "failure_conditions": [
                "restored from backup immediately",
                "ignored the suspicious email",
                "reset all passwords without investigation",
                "blamed the system admin"
            ]
        }
    ],
    "Adventure": [
        {
            "title": "The Lost Temple",
            "initial_setup": "Deep in the Amazon rainforest, you've discovered the entrance to an ancient temple.",
            "character1_name": "Dr. Rivera",
            "character2_name": "Guide Santos",
            "steps": [
                "Ancient markings warn of a curse",
                "You find a mechanism with multiple levers",
                "A strange humming sound emanates from deeper within",
                "The floor tiles show a peculiar pattern",
                "A beam of light reveals a hidden chamber"
            ],
            "success_conditions": [
                "documented the markings",
                "observed the pattern",
                "tested the mechanism carefully",
                "followed the light beam"
            ],
            "failure_conditions": [
                "ignored the warnings",
                "pulled levers randomly",
                "split up the group",
                "took artifacts without examination"
            ]
        }
    ],
    "Sci-Fi": [
        {
            "title": "The Quantum Anomaly",
            "initial_setup": "At a cutting-edge research facility, a quantum experiment has created an unexplained phenomenon.",
            "character1_name": "Dr. Zhang",
            "character2_name": "Engineer Parker",
            "steps": [
                "Quantum readings are off the charts",
                "Equipment starts behaving erratically",
                "A shimmer appears in the air",
                "Time seems to flow differently near the anomaly",
                "Multiple reality signatures detected"
            ],
            "success_conditions": [
                "monitored quantum fluctuations",
                "calibrated equipment",
                "documented time discrepancies",
                "maintained safe distance"
            ],
            "failure_conditions": [
                "shut down power immediately",
                "entered the anomaly",
                "ignored safety protocols",
                "attempted to contain without data"
            ]
        }
    ]
}

# Flatten categories for easy access by title
STORY_LOOKUP = {
    story["title"]: story 
    for category in STORY_CATEGORIES.values() 
    for story in category
}

# Character Response Templates
CHARACTER1_TEMPLATE = """
Context: You are {character1_name} in this story.
Story Progress: {story_log}
User's Latest Action: {user_input}

Respond to the user's action in character, considering:
1. Your role and personality
2. The current story situation
3. The potential consequences of their action

Response:
"""

CHARACTER2_TEMPLATE = """
Context: You are {character2_name} in this story.
Story Progress: {story_log}
User's Latest Action: {user_input}
Other Character's Response: {character1_response}

Respond to both the user and {character1_name}, considering:
1. Your role and personality
2. The current story developments
3. Your relationship with {character1_name}
4. The potential impact on the story's outcome

Response:
"""

def initialize_session_state():
    if 'story_state' not in st.session_state:
        st.session_state.story_state = StoryState()
    if 'selected_category' not in st.session_state:
        st.session_state.selected_category = list(STORY_CATEGORIES.keys())[0]
    if 'current_story' not in st.session_state:
        # Select random story from current category
        st.session_state.current_story = random.choice(STORY_CATEGORIES[st.session_state.selected_category])

def evaluate_outcome(state: StoryState, story: MicroStory) -> str:
    user_actions = " ".join(state.user_inputs).lower()
    
    # Count matches for success and failure conditions
    success_matches = sum(1 for cond in story["success_conditions"] if cond.lower() in user_actions)
    failure_matches = sum(1 for cond in story["failure_conditions"] if cond.lower() in user_actions)
    
    # Calculate success ratio
    success_ratio = success_matches / len(story["success_conditions"])
    
    if failure_matches >= 2:
        return "The story ends in failure. Critical mistakes were made."
    elif success_ratio >= 0.7:
        return "The story concludes successfully! Well done!"
    else:
        return "The story ends with mixed results. Some opportunities were missed."

def update_story_state(state: StoryState, user_input: str, char1_response: str, char2_response: str):
    state.current_step += 1
    state.user_inputs.append(user_input)
    state.character1_responses.append(char1_response)
    state.character2_responses.append(char2_response)
    
    # Check if story should end
    if state.current_step >= state.max_steps:
        state.story_outcome = evaluate_outcome(state, st.session_state.current_story)

def generate_character_response(
    character_name: str,
    story_log: List[str],
    user_input: str,
    other_response: Optional[str] = None,
    is_character1: bool = True
) -> str:
    # This would normally use watsonx.ai or another LLM
    # For now, return placeholder responses
    if is_character1:
        return f"{character_name}: That's an interesting approach. Let's see where this leads..."
    else:
        return f"{character_name}: I have my doubts about this, but we'll see..."

def main():
    st.set_page_config(page_title="Interactive Story", layout="wide")
    initialize_session_state()
    
    # Sidebar Configuration
    st.sidebar.header('Story Selection')
    st.sidebar.divider()
    
    # Category Selection
    selected_category = st.sidebar.selectbox(
        "Select Story Category",
        list(STORY_CATEGORIES.keys()),
        index=list(STORY_CATEGORIES.keys()).index(st.session_state.selected_category)
    )
    
    # Update category and story if changed
    if selected_category != st.session_state.selected_category:
        st.session_state.selected_category = selected_category
        st.session_state.current_story = random.choice(STORY_CATEGORIES[selected_category])
        st.session_state.story_state = StoryState()  # Reset state for new story
        st.rerun()
    
    # Display available stories in category
    with st.sidebar.expander(f"Available {selected_category} Stories"):
        for story in STORY_CATEGORIES[selected_category]:
            st.write(f"📖 {story['title']}")
    
    # Optional: Select specific story
    specific_story = st.sidebar.selectbox(
        "Select Specific Story",
        [story["title"] for story in STORY_CATEGORIES[selected_category]],
        index=[story["title"] for story in STORY_CATEGORIES[selected_category]].index(st.session_state.current_story["title"])
    )
    
    # Update if specific story changed
    if specific_story != st.session_state.current_story["title"]:
        st.session_state.current_story = STORY_LOOKUP[specific_story]
        st.session_state.story_state = StoryState()  # Reset state for new story
        st.rerun()
    
    # Display story stats
    st.sidebar.divider()
    st.sidebar.subheader("Story Progress")
    progress = (st.session_state.story_state.current_step / 5) * 100
    st.sidebar.progress(progress)
    st.sidebar.write(f"Step {st.session_state.story_state.current_step + 1}/5")
    
    # Create three columns
    col1, col2, col3 = st.columns(3)
    
    # Story Progress Column
    with col1:
        st.header("Story Progress")
        st.write(f"**{st.session_state.current_story['title']}**")
        st.write(st.session_state.current_story['initial_setup'])
        
        # Display story log
        for step_num, (step, user_input) in enumerate(zip(
            st.session_state.current_story['steps'][:st.session_state.story_state.current_step],
            st.session_state.story_state.user_inputs
        )):
            st.write(f"Step {step_num + 1}: {step}")
            st.write(f"Your action: {user_input}")
            st.write("---")
        
        # Display current step if story isn't finished
        if st.session_state.story_state.story_outcome is None:
            current_step = st.session_state.current_story['steps'][st.session_state.story_state.current_step]
            st.write(f"Current Situation: {current_step}")
    
    # Character 1 Column
    with col2:
        st.header(st.session_state.current_story['character1_name'])
        for response in st.session_state.story_state.character1_responses:
            st.write(response)
    
    # Character 2 Column
    with col3:
        st.header(st.session_state.current_story['character2_name'])
        for response in st.session_state.story_state.character2_responses:
            st.write(response)
    
    # User Input Section
    if st.session_state.story_state.story_outcome is None:
        user_input = st.text_input(
            "What do you do?",
            key=f"user_input_{st.session_state.story_state.current_step}"
        )
        
        if user_input:
            # Generate character responses
            char1_response = generate_character_response(
                st.session_state.current_story['character1_name'],
                st.session_state.story_state.story_log,
                user_input
            )
            
            char2_response = generate_character_response(
                st.session_state.current_story['character2_name'],
                st.session_state.story_state.story_log,
                user_input,
                char1_response,
                False
            )
            
            # Update state
            update_story_state(st.session_state.story_state, user_input, char1_response, char2_response)
            st.rerun()
    else:
        # Display story outcome
        st.write(st.session_state.story_state.story_outcome)
        if st.button("Start New Story"):
            st.session_state.story_state = StoryState()
            st.session_state.current_story = random.choice(MICRO_STORIES)
            st.rerun()

if __name__ == "__main__":
    main()