from typing import List, Tuple, Dict, Any, Optional
import logging
import re
from data_models import IdeaForm, IDEA_STAGES
from config import DEFAULT_SYSTEM_PROMPT, STAGES
from utils import (
    get_llm_response, extract_form_data, save_idea_to_database,
    load_idea_from_database, update_idea_in_database,
    get_db, clear_database, init_db, create_tables,
    perform_web_search, optimize_search_query,
    SessionLocal, InnovativeIdea
)

class InnovativeIdeaChatbot:
    def __init__(self):
        create_tables()
        init_db()
        self.idea_form = IdeaForm()
        self.chat_history = []
        self.idea_id = None
        self.current_stage = None
        self.api_key = None
        self.add_system_message(self.get_initial_greeting())

    def get_initial_greeting(self) -> str:
        greeting = """
Welcome to the Innovative Idea Generator! I'm Myamoto, your AI assistant designed to help you refine and develop your innovative ideas.
Here's how we'll work together:
1. We'll go through 10 stages to explore different aspects of your idea.
2. At each stage, I'll ask you questions and provide feedback to help you think deeper about your concept.
3. You can ask me questions at any time or request more information on a topic.
4. If you want to perform a web search for additional information, just start your message with '@' followed by your search query.
5. When you're ready to move to the next stage, simply type 'next'.
Let's start by exploring your innovative idea! What's the name of your idea, or would you like help coming up with one?
"""
        self.greeted = True
        return greeting

    def add_system_message(self, message: str):
        self.chat_history.append(("System", message))

    def set_api_key(self, api_key: str):
        self.api_key = api_key

    def activate_stage(self, stage_name: str) -> Optional[str]:
        self.current_stage = stage_name
        for stage in STAGES:
            if stage["name"] == stage_name:
                return f"Let's work on the '{stage_name}' stage. {stage['question']}"
        return None

    def process_stage_input(self, stage_name: str, message: str, model: str, system_prompt: str, thinking_budget: int) -> Tuple[List[Tuple[str, str]], Dict[str, Any]]:
        if self.current_stage != stage_name:
            activation_message = self.activate_stage(stage_name)
            if activation_message is None:
                error_message = f"Error: Unable to activate stage '{stage_name}'. Please check if the stage name is correct."
                self.chat_history.append(("System", error_message))
                return self.chat_history, self.idea_form.dict()
            self.chat_history.append(("System", activation_message))

        # Check for web search request
        if message.startswith('@'):
            search_query = message[1:].strip()
            optimized_query = optimize_search_query(search_query, model)
            search_results = perform_web_search(optimized_query)
            self.chat_history.append(("Human", message))
            self.chat_history.append(("AI", f"Here are the search results for '{optimized_query}':\n\n{search_results}"))
            return self.chat_history, self.idea_form.dict()

        # Generate the prompt for the current stage
        stage_prompt = self.generate_prompt_for_stage(stage_name)

        # Use the DEFAULT_SYSTEM_PROMPT from config.py
        formatted_system_prompt = DEFAULT_SYSTEM_PROMPT.format(
            current_stage=stage_name,
            stage_prompt=stage_prompt
        )

        # Combine the formatted system prompt and user's input
        combined_prompt = f"{formatted_system_prompt}\n\nUser input: {message}"

        # Get LLM response
        llm_response = get_llm_response(combined_prompt, model, thinking_budget, self.api_key)

        # Parse the LLM response to extract only the user-facing content
        parsed_response = self.parse_llm_response(llm_response)

        # Add the interaction to chat history
        self.chat_history.append(("Human", message))
        self.chat_history.append(("AI", parsed_response))

        # Extract form data from the LLM response
        form_data = extract_form_data(llm_response)
        
        # Update the idea form
        if stage_name in form_data:
            setattr(self.idea_form, stage_name.lower().replace(" ", "_"), form_data[stage_name])

        return self.chat_history, self.idea_form.dict()

    def parse_llm_response(self, response: str) -> str:
        # Remove content within <form_data> tags
        response = re.sub(r'<form_data>.*?</form_data>', '', response, flags=re.DOTALL)
        
        # Remove content within <reflection> tags
        response = re.sub(r'<reflection>.*?</reflection>', '', response, flags=re.DOTALL)
        
        # Remove content within <analysis> tags
        response = re.sub(r'<analysis>.*?</analysis>', '', response, flags=re.DOTALL)
        
        # Remove content within <summary> tags
        response = re.sub(r'<summary>.*?</summary>', '', response, flags=re.DOTALL)
        
        # Remove content within <step> tags
        response = re.sub(r'<step>.*?</step>', '', response, flags=re.DOTALL)
        
        # Remove any remaining HTML-like tags
        response = re.sub(r'<[^>]+>', '', response)
        
        # Remove extra whitespace and newlines
        response = re.sub(r'\s+', ' ', response).strip()
        
        return response

    def fill_out_form(self, current_stage: str, model: str, thinking_budget: int) -> Dict[str, str]:
        form_data = {}
        for stage in STAGES:
            stage_name = stage["name"]
            if stage_name == current_stage:
                # Generate new data for the current stage
                form_data[stage["field"]] = self.generate_form_data(stage_name, model, thinking_budget)
            else:
                # Use existing data for other stages
                form_data[stage["field"]] = getattr(self.idea_form, stage["field"], "")
        
        # Update the idea form
        for stage in STAGES:
            setattr(self.idea_form, stage["field"], form_data[stage["field"]])
        
        # Save to database
        try:
            new_session = SessionLocal()
            if self.idea_id:
                update_idea_in_database(self.idea_id, self.idea_form, new_session)
            else:
                self.idea_id = save_idea_to_database(self.idea_form, new_session)
            new_session.commit()
        except Exception as e:
            logging.error(f"Error saving idea to database: {str(e)}")
            new_session.rollback()
        finally:
            new_session.close()

        return form_data

    def generate_prompt_for_stage(self, stage: str) -> str:
        for s in IDEA_STAGES:
            if s.name == stage:
                return f"We are currently working on the '{stage}' stage. {s.question}"
        return f"We are currently working on the '{stage}' stage. Please provide relevant information."

    def reset(self):
        self.chat_history = []
        self.idea_form = IdeaForm()
        self.idea_id = None
        self.current_stage = None
        self.add_system_message(self.get_initial_greeting())
        try:
            new_session = SessionLocal()
            clear_database(new_session)
            new_session.commit()
        except Exception as e:
            logging.error(f"Error clearing database: {str(e)}")
            new_session.rollback()
        finally:
            new_session.close()
        return self.chat_history, self.idea_form.dict()

    def start_over(self):
        self.chat_history = []
        self.idea_form = IdeaForm()
        self.current_stage = None
        self.add_system_message(self.get_initial_greeting())
        
        try:
            new_session = SessionLocal()
            # Clear the existing database
            clear_database(new_session)
            
            # Create a new empty idea
            new_idea = InnovativeIdea()
            new_session.add(new_idea)
            new_session.commit()
            new_session.refresh(new_idea)
            
            # Update the idea_id
            self.idea_id = new_idea.id
            
            new_session.close()
        except Exception as e:
            logging.error(f"Error in start_over: {str(e)}")
            if 'new_session' in locals():
                new_session.rollback()
                new_session.close()
        
        return self.chat_history, self.idea_form.dict(), STAGES[0]["name"]

    def update_idea_form(self, stage_name: str, form_data: str):
        setattr(self.idea_form, stage_name.lower().replace(" ", "_"), form_data)
        try:
            new_session = SessionLocal()
            if self.idea_id:
                update_idea_in_database(self.idea_id, self.idea_form, new_session)
            else:
                self.idea_id = save_idea_to_database(self.idea_form, new_session)
            new_session.commit()
        except Exception as e:
            logging.error(f"Error updating idea form: {str(e)}")
            new_session.rollback()
        finally:
            new_session.close()

    def generate_form_data(self, stage: str, model: str, thinking_budget: int) -> str:
        # Prepare the conversation history for the LLM
        conversation = "\n".join([f"{role}: {message}" for role, message in self.chat_history])
        
        stage_prompt = self.generate_prompt_for_stage(stage)
        
        formatted_system_prompt = DEFAULT_SYSTEM_PROMPT.format(
            current_stage=stage,
            stage_prompt=stage_prompt
        )

        prompt = f"""
        {formatted_system_prompt}

        Based on the following conversation, extract the relevant information for the '{stage}' stage of the innovative idea:

        {conversation}

        Please provide a concise summary for the '{stage}' stage, focusing only on the information relevant to this stage. 
        Your response should be structured as follows:

        1. A brief analysis of the conversation related to this stage.
        2. A concise summary of the key points relevant to this stage.
        3. A suggested form entry for this stage, enclosed in <form_data></form_data> tags.
           The form entry should be in the format: "{stage}: Content"

        Remember to keep the form entry concise and directly related to the '{stage}' stage. Do not include information from other stages in the form entry.
        """

        # Get LLM response
        llm_response = get_llm_response(prompt, model, thinking_budget, self.api_key)

        # Extract form data from the LLM response
        form_data = extract_form_data(llm_response)
        
        return form_data.get(stage, "")

    def process_stage_input_stream(self, stage_name: str, message: str, model: str, system_prompt: str, thinking_budget: int):
        if self.current_stage != stage_name:
            activation_message = self.activate_stage(stage_name)
            if activation_message is None:
                error_message = f"Error: Unable to activate stage '{stage_name}'. Please check if the stage name is correct."
                self.chat_history.append(("System", error_message))
                yield self.chat_history, self.idea_form.dict()
                return

            self.chat_history.append(("System", activation_message))

        # Check for web search request
        if message.startswith('@'):
            search_query = message[1:].strip()
            optimized_query = optimize_search_query(search_query, model)
            search_results = perform_web_search(optimized_query)
            self.chat_history.append(("Human", message))
            self.chat_history.append(("AI", f"Here are the search results for '{optimized_query}':\n\n{search_results}"))
            yield self.chat_history, self.idea_form.dict()
            return

        # Generate the prompt for the current stage
        stage_prompt = self.generate_prompt_for_stage(stage_name)

        formatted_system_prompt = DEFAULT_SYSTEM_PROMPT.format(
            current_stage=stage_name,
            stage_prompt=stage_prompt
        )

        combined_prompt = f"{formatted_system_prompt}\n\nUser input: {message}"

        # Get LLM response
        llm_response = get_llm_response(combined_prompt, model, thinking_budget, self.api_key)
        parsed_response = self.parse_llm_response(llm_response)
        
        self.chat_history.append(("Human", message))
        self.chat_history.append(("AI", parsed_response))

        form_data = extract_form_data(llm_response)
        
        if stage_name in form_data:
            setattr(self.idea_form, stage_name.lower().replace(" ", "_"), form_data[stage_name])

        yield self.chat_history, self.idea_form.dict()

    def fill_out_form_stream(self, current_stage: str, model: str, thinking_budget: int):
        form_data = {}
        for stage in IDEA_STAGES:
            stage_name = stage.name
            if stage_name == current_stage:
                form_data[stage_name] = self.generate_form_data(stage_name, model, thinking_budget)
            else:
                form_data[stage_name] = getattr(self.idea_form, stage.field, "")
            yield form_data

        # Update the idea form
        for stage in IDEA_STAGES:
            setattr(self.idea_form, stage.field, form_data[stage.name])
        
        # Save to database
        try:
            new_session = SessionLocal()
            if self.idea_id:
                update_idea_in_database(self.idea_id, self.idea_form, new_session)
            else:
                self.idea_id = save_idea_to_database(self.idea_form, new_session)
            new_session.commit()
        except Exception as e:
            logging.error(f"Error saving idea to database: {str(e)}")
            new_session.rollback()
        finally:
            new_session.close()

    def generate_form_data_stream(self, stage: str, model: str, thinking_budget: int):
        conversation = "\n".join([f"{role}: {message}" for role, message in self.chat_history])
        
        stage_prompt = self.generate_prompt_for_stage(stage)
        
        formatted_system_prompt = DEFAULT_SYSTEM_PROMPT.format(
            current_stage=stage,
            stage_prompt=stage_prompt
        )

        prompt = f"""
        {formatted_system_prompt}

        Based on the following conversation, extract the relevant information for the '{stage}' stage of the innovative idea:

        {conversation}

        Please provide a concise summary for the '{stage}' stage, focusing only on the information relevant to this stage. 
        Your response should be structured as follows:

        1. A brief analysis of the conversation related to this stage.
        2. A concise summary of the key points relevant to this stage.
        3. A suggested form entry for this stage, enclosed in <form_data></form_data> tags.
           The form entry should be in the format: "{stage}: Content"

        Remember to keep the form entry concise and directly related to the '{stage}' stage. Do not include information from other stages in the form entry.
        """

        llm_response = get_llm_response(prompt, model, thinking_budget, self.api_key)
        form_data = extract_form_data(llm_response)
        return form_data.get(stage, "")

    def update_form_field(self, stage_name: str, value: str):
        field_name = stage_name.lower().replace(" ", "_")
        if field_name == 'team_roles':
            value = value.split(',')  # Convert string to list for team_roles
        setattr(self.idea_form, field_name, value)
        try:
            # Create a new session for this operation
            new_session = SessionLocal()
            if self.idea_id:
                update_idea_in_database(self.idea_id, self.idea_form, new_session)
            else:
                self.idea_id = save_idea_to_database(self.idea_form, new_session)
            new_session.commit()
        except Exception as e:
            logging.error(f"Error updating form field: {str(e)}")
            # If an error occurs, rollback the new session
            new_session.rollback()
        finally:
            # Always close the new session
            new_session.close()
        return value

# Add this new method to handle the "Fill out form" button click
def fill_out_form_button(chatbot: InnovativeIdeaChatbot, current_stage: str, model: str, thinking_budget: int):
    form_data = chatbot.fill_out_form(current_stage, model, thinking_budget)
    return {stage["field"]: form_data[stage["field"]] for stage in STAGES}