app.py

# -*- coding: utf-8 -*-
import streamlit as st
import requests
from openai import OpenAI
import google.generativeai as genai
import anthropic
import streamlit.components.v1 as components
import re
# from supabase import create_client # Keep commented if not using history/auth
import base64
import time
import json # For potentially structured Q&A

# ---------- Helper: Safe Rerun ----------
def safe_rerun():
    """Safely trigger a Streamlit rerun if the function exists."""
    if hasattr(st, "experimental_rerun"):
        st.experimental_rerun()
    elif hasattr(st, "rerun"):
        st.rerun()
    else:
        st.warning("Rerun function not available. Please update Streamlit.")

# ---------- Setup & API Client Initialization ----------
# (Keep your existing API key loading logic here - unchanged)
openai_client = None
genai_client = None
deepseek_api_key = None
claude_client = None
secrets_available = {"openai": False, "gemini": False, "deepseek": False, "claude": False}
secret_errors = []

# (Keep your existing API key loading and client initialization logic here)
# ... (Omitted for brevity, assume it's the same as before) ...
# OpenAI API Key
try:
    openai_api_key = st.secrets.get("OPENAI_API_KEY")
    if openai_api_key:
        openai_client = OpenAI(api_key=openai_api_key)
        secrets_available["openai"] = True
    else:
        secret_errors.append("Streamlit Secret `OPENAI_API_KEY` not found.")
except KeyError:
    secret_errors.append("Streamlit Secret `OPENAI_API_KEY` not found.")
except Exception as e:
    secret_errors.append(f"Error initializing OpenAI client: {e}")

# Gemini API Key (Google GenAI)
try:
    gemini_api_key = st.secrets.get("GEMINI_API_KEY")
    if gemini_api_key:
        genai.configure(api_key=gemini_api_key)
        genai_client = genai
        secrets_available["gemini"] = True
    else:
        secret_errors.append("Streamlit Secret `GEMINI_API_KEY` not found.")
except KeyError:
    secret_errors.append("Streamlit Secret `GEMINI_API_KEY` not found.")
except Exception as e:
    secret_errors.append(f"Error initializing Google GenAI client: {e}")

# DeepSeek API Key
try:
    deepseek_api_key = st.secrets.get("DEEPSEEK_API_KEY")
    if deepseek_api_key:
        secrets_available["deepseek"] = True
    else:
        secret_errors.append("Streamlit Secret `DEEPSEEK_API_KEY` not found.")
except KeyError:
    secret_errors.append("Streamlit Secret `DEEPSEEK_API_KEY` not found.")
except Exception as e:
    secret_errors.append(f"Error reading DeepSeek API key: {e}")

# CLAUDE API Key and Client Initialization
try:
    claude_api_key = st.secrets.get("CLAUDE_API_KEY")
    if claude_api_key:
        claude_client = anthropic.Anthropic(api_key=claude_api_key)
        secrets_available["claude"] = True
    else:
        secret_errors.append("Streamlit Secret `CLAUDE_API_KEY` not found.")
except KeyError:
    secret_errors.append("Streamlit Secret `CLAUDE_API_KEY` not found.")
except Exception as e:
    secret_errors.append(f"Error initializing Claude client: {e}")

any_secret_loaded = any(secrets_available.values())


# ---------- Model Configuration ----------
# (Keep your existing SUPPORTED_MODELS dictionary population logic here - unchanged)
# ... (Omitted for brevity, assume it's the same as before) ...
SUPPORTED_MODELS = {}

# OpenAI Models
if secrets_available["openai"] and openai_client:
    SUPPORTED_MODELS.update({
        "GPT-4o (OpenAI)": {"id": "gpt-4o", "provider": "openai", "client": openai_client},
        "GPT-4o Mini (OpenAI)": {"id": "gpt-4o-mini", "provider": "openai", "client": openai_client},
        "GPT-4 Turbo (OpenAI)": {"id": "gpt-4-turbo", "provider": "openai", "client": openai_client},
        "GPT-4 (OpenAI)": {"id": "gpt-4", "provider": "openai", "client": openai_client},
        "GPT-3.5 Turbo (OpenAI)": {"id": "gpt-3.5-turbo", "provider": "openai", "client": openai_client},
    })

# Gemini Models
if secrets_available["gemini"] and genai_client:
    SUPPORTED_MODELS.update({
        "Gemini 1.5 Pro (Google)": {"id": "gemini-1.5-pro-latest", "provider": "gemini", "client": genai_client},
        "Gemini 1.5 Flash (Google)": {"id": "gemini-1.5-flash-latest", "provider": "gemini", "client": genai_client},
        "Gemini 1.0 Pro (Google)": {"id": "gemini-1.0-pro", "provider": "gemini", "client": genai_client},
    })

# DeepSeek Models
if secrets_available["deepseek"] and deepseek_api_key:
    SUPPORTED_MODELS.update({
        "DeepSeek Chat": {"id": "deepseek-chat", "provider": "deepseek", "client": None},
        "DeepSeek Coder": {"id": "deepseek-coder", "provider": "deepseek", "client": None},
    })

# Claude Models
if secrets_available["claude"] and claude_client:
    SUPPORTED_MODELS.update({
        "Claude 3.5 Sonnet (Anthropic)": {"id": "claude-3-5-sonnet-20240620", "provider": "claude", "client": claude_client},
        "Claude 3 Haiku (Anthropic)": {"id": "claude-3-haiku-20240307", "provider": "claude", "client": claude_client},
        "Claude 3 Opus (Anthropic)": {"id": "claude-3-opus-20240229", "provider": "claude", "client": claude_client},
    })

# Determine default model based on preference and availability
DEFAULT_MODEL_PREFERENCE = [
    "GPT-4o Mini (OpenAI)",
    "Gemini 1.5 Flash (Google)",
    "Claude 3 Haiku (Anthropic)",
    "DeepSeek Chat",
    "GPT-3.5 Turbo (OpenAI)",
]
DEFAULT_MODEL = next((m for m in DEFAULT_MODEL_PREFERENCE if m in SUPPORTED_MODELS), None)
if not DEFAULT_MODEL and SUPPORTED_MODELS:
    DEFAULT_MODEL = next(iter(SUPPORTED_MODELS)) # Fallback to the first available


# ---------- Helper Functions for Generation ----------
# (Keep your existing _generate_with_... provider functions here - unchanged)
# ... (Omitted for brevity, assume they are the same as before) ...
def _generate_with_openai_provider(client, model_id, prompt, max_tokens, system_message=None):
    messages = []
    if system_message:
        messages.append({"role": "system", "content": system_message})
    messages.append({"role": "user", "content": prompt})
    try:
        response = client.chat.completions.create(
            model=model_id,
            messages=messages,
            temperature=0.6,
            max_tokens=max_tokens
        )
        return response.choices[0].message.content
    except Exception as e:
        st.error(f"❌ OpenAI API Error ({model_id}): {e}")
        return f"Error: OpenAI API call failed for {model_id}. Details: {e}"

def _generate_with_gemini_provider(client, model_id, prompt, max_tokens, system_message=None):
    full_prompt = f"{system_message}\n\n{prompt}" if system_message else prompt
    try:
        model = client.GenerativeModel(
            model_id,
            safety_settings={
                'HARM_CATEGORY_HARASSMENT': 'block_none',
                'HARM_CATEGORY_HATE_SPEECH': 'block_none',
                'HARM_CATEGORY_SEXUALLY_EXPLICIT': 'block_none',
                'HARM_CATEGORY_DANGEROUS_CONTENT': 'block_none',
            },
            generation_config=client.types.GenerationConfig(temperature=0.7)
        )
        response = model.generate_content(full_prompt)

        if response.parts:
             return "".join(part.text for part in response.parts if hasattr(part, 'text'))
        elif hasattr(response, 'text') and response.text:
             return response.text
        elif response.prompt_feedback.block_reason:
            reason = response.prompt_feedback.block_reason
            st.warning(f"Gemini response blocked ({model_id}). Reason: {reason}")
            return f"Error: Response blocked by API safety filters ({model_id}): {reason}"
        else:
            if response.candidates and response.candidates[0].finish_reason != "STOP":
                 st.warning(f"Gemini generation stopped unexpectedly ({model_id}). Reason: {response.candidates[0].finish_reason}")
                 return f"Error: Generation stopped unexpectedly ({model_id}). Reason: {response.candidates[0].finish_reason}"
            else:
                 st.warning(f"Gemini returned an empty or unexpected response ({model_id}).")
                 return f"Error: Gemini returned an empty response for {model_id}."

    except Exception as e:
        st.error(f"❌ Gemini SDK error ({model_id}): {e}")
        error_detail = getattr(e, 'message', str(e))
        if "API key not valid" in error_detail:
            return f"Error: Invalid Gemini API Key ({model_id}). Please check your Streamlit secrets."
        return f"Error: Gemini SDK call failed for {model_id}. Details: {error_detail}"


def _generate_with_deepseek_provider(api_key, model_id, prompt, max_tokens, system_message=None):
    headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
    messages = []
    if system_message:
        messages.append({"role": "system", "content": system_message})
    messages.append({"role": "user", "content": prompt})
    payload = {
        "model": model_id,
        "messages": messages,
        "temperature": 0.6,
        "max_tokens": max_tokens
    }
    try:
        response = requests.post("https://api.deepseek.com/chat/completions", headers=headers, json=payload, timeout=90)
        response.raise_for_status()
        response_data = response.json()
        if ("choices" in response_data and response_data["choices"] and
            "message" in response_data["choices"][0] and
            "content" in response_data["choices"][0]["message"]):
            return response_data["choices"][0]["message"]["content"]
        else:
            st.warning(f"DeepSeek returned an unexpected response structure ({model_id}): {response_data}")
            return f"Error: DeepSeek returned an unexpected structure for {model_id}."
    except requests.exceptions.RequestException as e:
        st.error(f"❌ DeepSeek API Request Error ({model_id}): {e}")
        return f"Error: DeepSeek API request failed for {model_id}. Details: {e}"
    except Exception as e:
        st.error(f"❌ DeepSeek Error processing response ({model_id}): {e}")
        return f"Error: DeepSeek processing failed for {model_id}. Details: {e}"

def _generate_with_claude_provider(client, model_id, prompt, max_tokens, system_message=None):
    try:
        message = client.messages.create(
            model=model_id,
            max_tokens=max_tokens,
            system=system_message if system_message else None,
            messages=[
                {"role": "user", "content": prompt}
            ]
        )
        content = ""
        if message.content:
           content = "\n".join([block.text for block in message.content if hasattr(block, "text")])
        return content
    except Exception as e:
        st.error(f"❌ Claude API Error ({model_id}): {e}")
        if isinstance(e, anthropic.AuthenticationError):
             return f"Error: Claude authentication failed ({model_id}). Check your API key."
        return f"Error: Claude API call failed for {model_id}. Details: {e}"

def generate_with_selected_model(selected_model_name, prompt, max_tokens=2000, system_message=None):
    """Generates text using the chosen model, handling provider specifics."""
    if not any_secret_loaded or not SUPPORTED_MODELS:
        st.error("Error: No API keys loaded or models available. Configure secrets.")
        return None

    if selected_model_name not in SUPPORTED_MODELS:
        st.error(f"Selected model '{selected_model_name}' is not configured or unavailable.")
        original_choice = selected_model_name
        selected_model_name = DEFAULT_MODEL
        if not selected_model_name:
            st.error("Fatal: Default model is also unavailable. Cannot proceed.")
            return None
        st.warning(f"Falling back from '{original_choice}' to default model: {DEFAULT_MODEL}")
        st.session_state.model_choice = DEFAULT_MODEL # Update state on fallback

    model_config = SUPPORTED_MODELS[selected_model_name]
    provider = model_config["provider"]
    model_id = model_config["id"]
    client = model_config.get("client")

    st.info(f"Generating with: **{selected_model_name}**")
    start_time = time.time()
    result = f"Error: Provider '{provider}' not implemented."

    try:
        if provider == "openai":
            if not client: result = f"Error: OpenAI client not initialized for {selected_model_name}."
            else: result = _generate_with_openai_provider(client, model_id, prompt, max_tokens, system_message)
        elif provider == "gemini":
            if not client: result = f"Error: Gemini client not initialized for {selected_model_name}."
            else: result = _generate_with_gemini_provider(client, model_id, prompt, max_tokens, system_message)
        elif provider == "deepseek":
            if not deepseek_api_key: result = f"Error: DeepSeek API key not available for {selected_model_name}."
            else: result = _generate_with_deepseek_provider(deepseek_api_key, model_id, prompt, max_tokens, system_message)
        elif provider == "claude":
            if not client: result = f"Error: Claude client not initialized for {selected_model_name}."
            else: result = _generate_with_claude_provider(client, model_id, prompt, max_tokens, system_message)
    except Exception as e:
        st.error(f"❌ Unexpected error during generation with {selected_model_name}: {e}")
        result = f"Error: Unexpected failure during {provider} generation. Details: {e}"

    end_time = time.time()
    duration = end_time - start_time
    # st.caption(f"Generation took {duration:.2f} seconds.") # Less verbose

    if isinstance(result, str) and result.startswith("Error:"):
        # Error already logged by provider function
        return None
    return result

# --- Mermaid Diagram Helper ---
# (Keep your existing Mermaid helper functions here - unchanged)
# ... (Omitted for brevity, assume they are the same as before) ...
def is_valid_mermaid(code):
    if not isinstance(code, str): return False
    code_lower = code.strip().lower()
    return bool(re.search(r"^\s*(%%.*?\n)*\s*(graph|flowchart|sequenceDiagram|classDiagram|stateDiagram|erDiagram|gantt|pie|gitGraph)", code_lower, re.MULTILINE))

def render_mermaid_diagram(mermaid_code, key):
    if not isinstance(mermaid_code, str) or not mermaid_code.strip():
        st.warning(f"Mermaid code is empty or invalid (Key: {key}).")
        return

    cleaned_code = re.sub(r"^```mermaid\s*\n?", "", mermaid_code, flags=re.IGNORECASE | re.MULTILINE).strip()
    cleaned_code = re.sub(r"\n?```\s*$", "", cleaned_code).strip()

    if not is_valid_mermaid(cleaned_code):
        st.warning(f"⚠️ Mermaid diagram might not render correctly (Key: {key}). Check syntax. Displaying raw code.")
        st.code(cleaned_code, language="mermaid")
        return

    container_id = f"mermaid-container-{key}"
    mermaid_id = f"mermaid-{key}"

    components.html(
        f"""
        <div id="{container_id}" style="background-color: white; padding: 10px; border-radius: 5px; overflow: auto;">
            <pre class="mermaid" id="{mermaid_id}">
            {cleaned_code}
            </pre>
        </div>
        <script type="module">
            try {{
                const mermaid = (await import('https://cdn.jsdelivr.net/npm/mermaid@10/dist/mermaid.esm.min.mjs')).default;
                mermaid.initialize({{ startOnLoad: false, theme: 'default' }});
                const checkElement = setInterval(() => {{
                    const el = document.getElementById('{mermaid_id}');
                    if (el) {{
                        clearInterval(checkElement);
                        mermaid.run({{ nodes: [el] }});
                    }}
                }}, 100);
                setTimeout(() => clearInterval(checkElement), 5000);
            }} catch (e) {{
                console.error("Mermaid rendering error (Key: {key}):", e);
                const container = document.getElementById('{container_id}');
                if(container) container.innerHTML = "<p style='color:red;'>Error rendering Mermaid diagram. Check browser console.</p>";
            }}
        </script>
        """,
        height=500, scrolling=True,
    )

# ---------- Initialize Session State (Provibe Workflow) ----------
# Core workflow steps
if 'current_step' not in st.session_state:
    st.session_state.current_step = "input_idea" # input_idea -> refine_idea -> review_idea -> generate_docs -> display_docs
if 'processing' not in st.session_state: # General flag for disabling buttons
    st.session_state.processing = False

# Input data
if 'initial_product_idea' not in st.session_state:
    st.session_state.initial_product_idea = ""
if 'tech_stack_hint' not in st.session_state:
    st.session_state.tech_stack_hint = ""
if 'model_choice' not in st.session_state:
    st.session_state.model_choice = DEFAULT_MODEL

# Refinement stage data (NEW states added here)
if 'refinement_sub_step' not in st.session_state:
    # Tracks progress within the 'refine_idea' step
    # Possible values: 'generate_questions', 'await_answers', 'generate_final_refinement'
    st.session_state.refinement_sub_step = 'generate_questions'
if 'clarifying_questions' not in st.session_state:
    # Stores the list of questions generated by the AI
    st.session_state.clarifying_questions = []
if 'user_answers' not in st.session_state:
    # Stores user answers, dictionary mapping question index to answer string
    st.session_state.user_answers = {}

# Output/Generated data
if 'refined_idea_content' not in st.session_state: # Stores the final AI-refined idea (after Q&A)
    st.session_state.refined_idea_content = None
if 'confirmed_idea_content' not in st.session_state: # Stores the user-confirmed/edited idea
    st.session_state.confirmed_idea_content = None
if 'selected_docs_to_generate' not in st.session_state: # Stores user selection for optional docs
    st.session_state.selected_docs_to_generate = {}
if 'generated_docs' not in st.session_state: # Stores content of generated optional docs
    st.session_state.generated_docs = {}


# ---------- Define Document Options (Align with Provibe Output) ----------
# (Keep your existing doc_options dictionary here - unchanged)
# ... (Omitted for brevity, assume it's the same as before, including PRD option) ...
doc_options = {
    "prd": {
        "label": "Product Requirements Document (PRD)",
        "prompt_func": lambda idea, hint: f"""
# --- PROMPT: Insert your specific PRD generation prompt here ---
# Example: Write a comprehensive Product Requirements Document (PRD) based strictly on the following confirmed product description. Include sections like Introduction, Goals, Target Audience, Features (with details), User Stories, Design Considerations, Non-Functional Requirements, Open Issues, and Future Considerations. Ensure the PRD is detailed, clear, and actionable for a development team.
# --- End PRD Prompt ---

**Confirmed Product Description:**
---
{idea}
---
**Optional Preferences/Hints (Consider if relevant):**
{hint if hint else "None provided"}
""",
        "system_message": "You are an expert Product Manager tasked with writing a detailed and professional PRD.",
        "max_tokens": 3500, # Allow more tokens for PRD
        "display_func": lambda content, key: st.markdown(content),
        "download_filename": "prd.md",
        "mime": "text/markdown",
    },
    "user_flow_text": {
        "label": "User Flow (Text Description)",
        "prompt_func": lambda idea, hint: f"""
# --- PROMPT: Insert your specific User Flow (Text) generation prompt here ---
# Example: Based on the product description below, outline the primary user flow step-by-step, from initial interaction to achieving the core goal. Describe each step clearly.
# --- End User Flow (Text) Prompt ---

**Product Description:**
---
{idea}
---
**Preferences/Hints:** {hint if hint else "None provided"}
""",
        "system_message": "You are a UX designer describing a key user journey.",
        "max_tokens": 1000,
        "display_func": lambda content, key: st.markdown(content),
        "download_filename": "user_flow.md",
        "mime": "text/markdown",
    },
    "user_flow_mermaid": {
        "label": "User Flow Diagram (Mermaid)",
        "prompt_func": lambda idea, hint: f"""
# --- PROMPT: Insert your specific User Flow (Mermaid) generation prompt here ---
# Example: Generate a Mermaid flowchart diagram representing the primary user flow for the product described below. Use standard flowchart syntax (graph TD, nodes, arrows). Ensure the diagram is clear and accurately reflects the user journey. Start the code block with ```mermaid and end it with ```. Do not include any other text before or after the code block.
# --- End User Flow (Mermaid) Prompt ---

**Product Description:**
---
{idea}
---
**Preferences/Hints:** {hint if hint else "None provided"}
""",
        "system_message": "You are an expert in creating Mermaid diagrams, specifically flowcharts for user journeys.",
        "max_tokens": 1000,
        "render_func": render_mermaid_diagram, # Special rendering
        "code_language": "mermaid",
        "download_filename": "user_flow_diagram.mmd",
        "mime": "text/plain",
    },
    "frontend_arch": {
        "label": "Frontend Architecture Notes",
        "prompt_func": lambda idea, hint: f"""
# --- PROMPT: Insert your Frontend Architecture prompt here ---
# Example: Based on the product description and hints, suggest a suitable frontend architecture. Describe key components, recommended libraries/frameworks (consider hints like 'React Native'), state management approach, and potential component breakdown.
# --- End Frontend Architecture Prompt ---

**Product Description:**
---
{idea}
---
**Preferences/Hints:** {hint if hint else "None provided"}
""",
        "system_message": "You are a frontend architect designing a web/mobile application.",
         "max_tokens": 1500,
        "display_func": lambda content, key: st.markdown(content),
        "download_filename": "frontend_architecture.md",
        "mime": "text/markdown",
    },
    "backend_arch": {
        "label": "Backend Architecture Notes",
        "prompt_func": lambda idea, hint: f"""
# --- PROMPT: Insert your Backend Architecture prompt here ---
# Example: Based on the product description and hints, propose a backend architecture. Discuss potential API design (e.g., RESTful), choice of language/framework, database considerations (type, scaling), authentication/authorization strategy, and key microservices or modules if applicable.
# --- End Backend Architecture Prompt ---

**Product Description:**
---
{idea}
---
**Preferences/Hints:** {hint if hint else "None provided"}
""",
        "system_message": "You are a backend/systems architect designing the server-side logic and infrastructure.",
         "max_tokens": 1500,
        "display_func": lambda content, key: st.markdown(content),
        "download_filename": "backend_architecture.md",
        "mime": "text/markdown",
    },
     "system_arch_mermaid": {
        "label": "System Architecture Diagram (Mermaid)",
        "prompt_func": lambda idea, hint: f"""
# --- PROMPT: Insert your System Architecture (Mermaid) prompt here ---
# Example: Generate a Mermaid diagram illustrating the high-level system architecture for the product described below. Include key components like frontend client, backend API, database, authentication service, and any major third-party integrations mentioned or implied. Use appropriate Mermaid diagram syntax (e.g., graph TD or C4 model elements if suitable). Start the code block with ```mermaid and end it with ```. Do not include any other text before or after the code block.
# --- End System Architecture (Mermaid) Prompt ---

**Product Description:**
---
{idea}
---
**Preferences/Hints:** {hint if hint else "None provided"}
""",
        "system_message": "You create system architecture diagrams using Mermaid syntax.",
        "max_tokens": 1000,
        "render_func": render_mermaid_diagram,
        "code_language": "mermaid",
        "download_filename": "system_architecture.mmd",
        "mime": "text/plain",
    },
    "db_schema": {
        "label": "Database Schema (SQL)",
        "prompt_func": lambda idea, hint: f"""
# --- PROMPT: Insert your Database Schema (SQL) prompt here ---
# Example: Based on the product description, design a preliminary relational database schema. Provide SQL `CREATE TABLE` statements for the primary entities, including relevant columns, data types, primary keys, and foreign key relationships. Assume a PostgreSQL syntax unless hints suggest otherwise.
# --- End Database Schema (SQL) Prompt ---

**Product Description:**
---
{idea}
---
**Preferences/Hints:** {hint if hint else "None provided"}
""",
        "system_message": "You are a database administrator designing a schema.",
        "max_tokens": 1500,
        "display_func": lambda content, key: st.code(content, language='sql'), # Use code block for SQL
        "code_language": "sql",
        "download_filename": "database_schema.sql",
        "mime": "text/x-sql",
    },
    "project_structure": {
        "label": "Project Folder Structure",
        "prompt_func": lambda idea, hint: f"""
# --- PROMPT: Insert your Project Structure prompt here ---
# Example: Suggest a logical file and folder structure for a project implementing the described product. Consider frontend, backend, shared components, tests, etc., based on the description and any tech stack hints. Present it as a simple tree structure.
# --- End Project Structure Prompt ---

**Product Description:**
---
{idea}
---
**Preferences/Hints:** {hint if hint else "None provided"}
""",
        "system_message": "You are suggesting a clean project layout for a software development team.",
        "max_tokens": 800,
        "display_func": lambda content, key: st.code(content, language='bash'), # Use code block for structure
        "code_language": "bash",
        "download_filename": "project_structure.txt",
        "mime": "text/plain",
    },
}


# ---------- UI Layout (Provibe Workflow) ----------
st.set_page_config(layout="wide", page_title="Provibe Prompt Tester")
st.title("🧪 Provibe Prompt Tester (with Q&A Refinement)")
st.caption("Test and refine prompts for the Provibe document generation workflow, including interactive Q&A.")

# Display API Key Errors
if secret_errors:
    st.error("API Key Configuration Issues:")
    for error in secret_errors:
        st.error(f"- {error}")
if not any_secret_loaded or not SUPPORTED_MODELS:
    st.error("No API keys loaded or LLM models available. Configure secrets.")
    st.stop()

# --- Workflow Steps ---

# ---------- Step 1: Input Initial Idea ----------
if st.session_state.current_step == "input_idea":
    st.header("Step 1: Input Product Idea")
    with st.form(key="idea_form"):
        initial_idea_input = st.text_area(
            "💡 Enter the initial product idea:", height=150,
            value=st.session_state.initial_product_idea,
            help="The raw concept or description."
        )
        tech_hint_input = st.text_input(
            "⚙️ Optional: Tech Stack Hints or Constraints",
            placeholder="e.g., Use React, target mobile, needs offline support",
            value=st.session_state.tech_stack_hint,
            help="Any preferences to guide AI generation."
        )
        available_model_names = list(SUPPORTED_MODELS.keys())
        default_model_key = st.session_state.get('model_choice', DEFAULT_MODEL)
        default_index = available_model_names.index(default_model_key) if default_model_key in available_model_names else 0
        model_choice_input = st.selectbox(
            "🧠 Choose AI model for all steps:",
            options=available_model_names,
            index=default_index,
            key="model_select",
            help="This model will be used for refinement and document generation."
        )
        submit_idea_button = st.form_submit_button(
            label="➡️ Start Interactive Refinement",
            use_container_width=True,
            disabled=st.session_state.processing
        )

    if submit_idea_button and initial_idea_input:
        st.session_state.initial_product_idea = initial_idea_input
        st.session_state.tech_stack_hint = tech_hint_input
        st.session_state.model_choice = model_choice_input
        # Reset states for the refinement process
        st.session_state.clarifying_questions = []
        st.session_state.user_answers = {}
        st.session_state.refined_idea_content = None
        st.session_state.confirmed_idea_content = None
        st.session_state.generated_docs = {}
        st.session_state.selected_docs_to_generate = {}
        # Set the next step and the initial sub-step for refinement
        st.session_state.current_step = "refine_idea"
        st.session_state.refinement_sub_step = "generate_questions"
        st.session_state.processing = True # Start processing
        safe_rerun()
    elif submit_idea_button:
        st.warning("Please enter a product idea.")


# ---------- Step 2: Interactive Refinement (Q&A) ----------
if st.session_state.current_step == "refine_idea":
    st.header("Step 2: Interactive Idea Refinement")

    # --- Sub-Step 2a: Generate Clarifying Questions ---
    if st.session_state.refinement_sub_step == "generate_questions":
        st.info(f"Using **{st.session_state.model_choice}** to generate clarifying questions. Please wait.")
        with st.spinner("AI is preparing questions..."):

            # --- PROMPT: Define the Question Generation Prompt ---
            question_gen_prompt = f"""
# --- PROMPT: Insert your Question Generation prompt here ---
# Example: Based on the initial product idea and hints below, generate 3-5 specific clarifying questions for the user. These questions should help elicit more detail about key features, target audience, technical constraints, or core functionality needed to write a better product specification. Output *only* the questions, each on a new line, starting with '- '. Do not include numbering or any other text.
# --- End Question Generation Prompt ---

**Initial Product Idea:**
---
{st.session_state.initial_product_idea}
---
**Optional Preferences/Hints Provided:**
{st.session_state.tech_stack_hint if st.session_state.tech_stack_hint else "None provided"}
"""
            # --- End Question Generation Prompt ---

            system_message_qa = "You are an AI assistant helping to clarify a product idea by asking relevant questions."
            max_tokens_qa = 300

            questions_raw = generate_with_selected_model(
                st.session_state.model_choice,
                question_gen_prompt,
                max_tokens=max_tokens_qa,
                system_message=system_message_qa
            )

            if questions_raw and not questions_raw.startswith("Error:"):
                # Parse the questions (assuming one question per line, maybe starting with '- ')
                st.session_state.clarifying_questions = [
                    q.strip('- ') for q in questions_raw.strip().split('\n') if q.strip() and q.strip() != '-'
                ]
                if st.session_state.clarifying_questions:
                    st.session_state.user_answers = {i: "" for i in range(len(st.session_state.clarifying_questions))} # Initialize empty answers
                    st.session_state.refinement_sub_step = "await_answers"
                else:
                     st.warning("AI generated questions but they seem empty or incorrectly formatted. Proceeding without Q&A.")
                     # Fallback: Skip Q&A and go directly to final refinement based only on initial idea
                     st.session_state.refinement_sub_step = "generate_final_refinement"
                     st.session_state.clarifying_questions = [] # Ensure it's empty
                     st.session_state.user_answers = {}

            else:
                st.error("Failed to generate clarifying questions. Check API errors or model selection.")
                st.warning("Proceeding to refine based only on the initial idea (skipping Q&A).")
                st.session_state.refinement_sub_step = "generate_final_refinement" # Skip Q&A on failure
                st.session_state.clarifying_questions = [] # Ensure it's empty
                st.session_state.user_answers = {}

            st.session_state.processing = False # Done generating questions (or failed)
            safe_rerun()

    # --- Sub-Step 2b: Display Questions and Collect Answers ---
    elif st.session_state.refinement_sub_step == "await_answers":
        st.info("Please answer the following questions to help refine the product idea:")
        with st.form("answers_form"):
            # Display generated questions and input fields for answers
            for i, question in enumerate(st.session_state.clarifying_questions):
                st.session_state.user_answers[i] = st.text_area(
                    f"❓ {question}",
                    key=f"answer_{i}",
                    value=st.session_state.user_answers.get(i, ""), # Preserve answers on rerun
                    height=100
                )

            submit_answers_button = st.form_submit_button(
                "➡️ Submit Answers & Generate Refined Description",
                use_container_width=True,
                disabled=st.session_state.processing
            )

        if submit_answers_button:
            # Basic check if answers are provided (optional)
            # if not all(st.session_state.user_answers.values()):
            #     st.warning("Please try to answer all questions for the best result.")
            # else:
            st.session_state.refinement_sub_step = "generate_final_refinement"
            st.session_state.processing = True # Start final refinement generation
            safe_rerun()

        # Option to go back
        if st.button("⬅️ Back to Idea Input (Discard Q&A)", disabled=st.session_state.processing):
            st.session_state.current_step = "input_idea"
            # Clear Q&A state
            st.session_state.clarifying_questions = []
            st.session_state.user_answers = {}
            safe_rerun()


    # --- Sub-Step 2c: Generate Final Refined Description (using Q&A) ---
    elif st.session_state.refinement_sub_step == "generate_final_refinement":
        st.info(f"Using **{st.session_state.model_choice}** to generate the final refined description based on the idea and your answers. Please wait.")
        with st.spinner("AI is synthesizing the refined description..."):

            # Prepare Q&A string for the prompt
            qa_summary = "\n".join([
                f"Q: {st.session_state.clarifying_questions[i]}\nA: {answer}"
                for i, answer in st.session_state.user_answers.items() if answer # Include only answered questions
            ]) if st.session_state.user_answers else "No questions were answered."


            # --- PROMPT: Define the Final Refinement Prompt (using Q&A) ---
            final_refinement_prompt = f"""
# --- PROMPT: Insert your Final Refinement prompt (using Q&A) here ---
# Example: Based on the initial product idea, user preferences, and the following question-answer pairs, generate a concise yet comprehensive 'Refined Product Description'. Synthesize all the information into a well-structured description covering the core value proposition, key features, target audience, and any clarified technical aspects. This description will be the basis for generating all subsequent documents.
# --- End Final Refinement Prompt ---

**Initial Product Idea:**
---
{st.session_state.initial_product_idea}
---
**Optional Preferences/Hints Provided:**
{st.session_state.tech_stack_hint if st.session_state.tech_stack_hint else "None provided"}
---
**Clarifying Questions & User Answers:**
---
{qa_summary}
---
"""
            # --- End Final Refinement Prompt ---

            system_message_final_refine = "You are an AI assistant synthesizing information into a final product specification."
            max_tokens_final_refine = 1500 # Allow slightly more tokens for synthesis

            final_refined_content = generate_with_selected_model(
                st.session_state.model_choice,
                final_refinement_prompt,
                max_tokens=max_tokens_final_refine,
                system_message=system_message_final_refine
            )

            if final_refined_content and not final_refined_content.startswith("Error:"):
                st.session_state.refined_idea_content = final_refined_content
                st.session_state.current_step = "review_idea" # Move to the next main step
            else:
                st.error("Failed to generate the final refined description.")
                # Option to retry or go back might be added here
                st.session_state.current_step = "input_idea" # Go back if failed

            st.session_state.processing = False # End processing
            safe_rerun()


# ---------- Step 3: Review and Confirm Final Idea ----------
if st.session_state.current_step == "review_idea":
    st.header("Step 3: Review and Confirm Final Refined Idea")
    if st.session_state.refined_idea_content:
        st.info("Review the AI's final refined description below (generated using your answers). Edit it as needed. This **final text** will be used to generate all documents.")

        # Display Q&A for context if available
        if st.session_state.clarifying_questions and st.session_state.user_answers:
             with st.expander("View Q&A used for this refinement"):
                 for i, q in enumerate(st.session_state.clarifying_questions):
                     st.markdown(f"**Q:** {q}")
                     st.markdown(f"**A:** {st.session_state.user_answers.get(i, '_No answer_')}")
                 st.markdown("---")


        edited_idea = st.text_area(
            "✏️ **Edit Final Refined Description:**",
            value=st.session_state.refined_idea_content,
            height=350,
            key="final_refined_idea_edit_area",
            help="Make any necessary corrections or additions."
        )

        button_col1, button_col2 = st.columns(2)
        with button_col1:
            confirm_button = st.button(
                "✅ Confirm & Proceed to Generate Docs",
                key="confirm_final_idea_button",
                use_container_width=True,
                disabled=st.session_state.processing
            )
        with button_col2:
             back_button = st.button(
                 "⬅️ Back to Idea Input (Start Over)",
                 key="back_to_input_final_button",
                 use_container_width=True,
                 disabled=st.session_state.processing
             )

        if confirm_button:
            if not edited_idea.strip():
                st.warning("The refined description cannot be empty.")
            else:
                st.session_state.confirmed_idea_content = edited_idea
                # Reset generation states
                st.session_state.generated_docs = {}
                st.session_state.selected_docs_to_generate = {k: False for k in doc_options} # Reset selections
                st.session_state.current_step = "generate_docs"
                safe_rerun()
        if back_button:
             st.session_state.current_step = "input_idea"
             # Clear refinement & Q&A state
             st.session_state.refined_idea_content = None
             st.session_state.clarifying_questions = []
             st.session_state.user_answers = {}
             safe_rerun()

    else:
        st.error("No refined idea content found. Please go back to Step 1.")
        if st.button("⬅️ Back to Idea Input"):
            st.session_state.current_step = "input_idea"
            safe_rerun()


# ---------- Step 4: Select and Generate Documents ----------
if st.session_state.current_step == "generate_docs":
    st.header("Step 4: Generate Product Documents")
    if st.session_state.confirmed_idea_content:
        st.markdown("**Based on this confirmed final description:**")
        with st.expander("View Confirmed Description", expanded=False):
             st.markdown(f"> {st.session_state.confirmed_idea_content}")

        st.subheader("Select Documents to Generate:")
        num_doc_options = len(doc_options)
        cols = st.columns(min(num_doc_options, 3))
        doc_keys = list(doc_options.keys())

        for i, key in enumerate(doc_keys):
            config = doc_options[key]
            with cols[i % 3]:
                if key not in st.session_state.selected_docs_to_generate:
                     st.session_state.selected_docs_to_generate[key] = False
                st.session_state.selected_docs_to_generate[key] = st.checkbox(
                    config["label"],
                    value=st.session_state.selected_docs_to_generate.get(key, False),
                    key=f"checkbox_{key}",
                    disabled=st.session_state.processing
                )

        generate_button = st.button(
            "🚀 Generate Selected Documents",
            key="generate_docs_button",
            use_container_width=True,
            disabled=st.session_state.processing
        )
        back_to_review_button = st.button(
             "⬅️ Back to Review Final Idea",
              key="back_to_review_final_button",
              use_container_width=True,
              disabled=st.session_state.processing
        )


        if generate_button:
            selected_keys = [k for k, v in st.session_state.selected_docs_to_generate.items() if v]
            if not selected_keys:
                st.warning("Please select at least one document type to generate.")
            else:
                st.session_state.processing = True
                st.session_state.generated_docs = {}
                st.info(f"⏳ Generating {len(selected_keys)} selected document(s) using {st.session_state.model_choice}...")
                progress_bar = st.progress(0)
                generation_successful = True

                for i, key in enumerate(selected_keys):
                    config = doc_options[key]
                    st.write(f"   - Generating {config['label']}...")
                    with st.spinner(f"AI processing {config['label']}..."):
                        prompt = config["prompt_func"](
                            st.session_state.confirmed_idea_content,
                            st.session_state.tech_stack_hint
                        )
                        system_msg = config.get("system_message")
                        max_tok = config.get("max_tokens", 2000)

                        content = generate_with_selected_model(
                            st.session_state.model_choice,
                            prompt,
                            max_tokens=max_tok,
                            system_message=system_msg
                        )

                        if content and not content.startswith("Error:"):
                            st.session_state.generated_docs[key] = content
                        else:
                            st.session_state.generated_docs[key] = f"Error: Failed to generate {config['label']}."
                            generation_successful = False
                            st.error(f"   - Failed to generate {config['label']}. See logs above.")
                    progress_bar.progress((i + 1) / len(selected_keys))
                    time.sleep(0.1)

                progress_bar.empty()
                st.session_state.processing = False

                if generation_successful:
                    st.success("✅ Document generation complete!")
                else:
                    st.warning("⚠️ Some documents could not be generated.")

                st.session_state.current_step = "display_docs"
                safe_rerun()

        if back_to_review_button:
            st.session_state.current_step = "review_idea"
            # Keep confirmed idea, but allow editing again
            # The refined_idea_content should still hold the content before editing
            safe_rerun()

    else:
        st.error("Confirmed idea content is missing. Please restart the process from Step 1.")
        if st.button("⬅️ Restart Process"):
            # Reset key states
            st.session_state.current_step = "input_idea"
            st.session_state.initial_product_idea = ""
            st.session_state.tech_stack_hint = ""
            st.session_state.refined_idea_content = None
            st.session_state.confirmed_idea_content = None
            st.session_state.clarifying_questions = []
            st.session_state.user_answers = {}
            st.session_state.generated_docs = {}
            st.session_state.selected_docs_to_generate = {}
            safe_rerun()


# ---------- Step 5: Display Generated Documents ----------
if st.session_state.current_step == "display_docs":
    st.header("Step 5: Generated Documents")

    if not st.session_state.generated_docs:
         st.info("No documents were generated in the previous step.")
    else:
        st.markdown("**Review the generated documents below:**")
        display_order = [key for key in doc_options if key in st.session_state.generated_docs]

        for key in display_order:
            content = st.session_state.generated_docs.get(key)
            if content:
                config = doc_options[key]
                st.subheader(f"📄 {config['label']}")
                is_error = isinstance(content, str) and content.startswith("Error:")

                if is_error:
                    st.error(content)
                else:
                    # Display/Render content
                    if config.get("render_func"):
                        try: config["render_func"](content, key=f"render_{key}")
                        except Exception as e: st.error(f"Render Error: {e}"); st.code(content)
                    elif config.get("display_func"):
                        try: config["display_func"](content, key=f"display_{key}")
                        except Exception as e: st.error(f"Display Error: {e}"); st.text(content)
                    else: st.markdown(content)

                    # Download button
                    try:
                        download_data = content.encode('utf-8') if isinstance(content, str) else str(content).encode('utf-8')
                        st.download_button(
                            label=f"📥 Download {config['label']}", data=download_data,
                            file_name=config["download_filename"], mime=config.get("mime", "text/plain"),
                            key=f"download_{key}"
                        )
                    except Exception as e: st.warning(f"Download Error: {e}")

                    # Show raw content option
                    if config.get("render_func") or config.get("code_language"):
                         if st.checkbox(f"🔍 Show raw content for {config['label']}", key=f"show_raw_{key}", value=False):
                              st.code(content, language=config.get("code_language", None))

                st.markdown("---")

    # Navigation buttons
    button_col1, button_col2 = st.columns(2)
    with button_col1:
        generate_more_button = st.button("🔄 Generate Different Documents", key="generate_more_button", use_container_width=True)
    with button_col2:
        restart_all_button = st.button("⏪ Start New Idea", key="restart_all_button", use_container_width=True)

    if generate_more_button:
         st.session_state.current_step = "generate_docs"
         st.session_state.generated_docs = {}
         safe_rerun()

    if restart_all_button:
        st.session_state.current_step = "input_idea"
        st.session_state.initial_product_idea = ""
        st.session_state.tech_stack_hint = ""
        st.session_state.refined_idea_content = None
        st.session_state.confirmed_idea_content = None
        st.session_state.clarifying_questions = []
        st.session_state.user_answers = {}
        st.session_state.generated_docs = {}
        st.session_state.selected_docs_to_generate = {}
        safe_rerun()


# ---------- Footer ----------
st.markdown("---")
footer_model_choice = st.session_state.get('model_choice', 'N/A')
st.caption(f"Using model: **{footer_model_choice}** | Workflow Step: **{st.session_state.get('current_step', 'N/A')}**"
           f"{' (Sub-step: ' + st.session_state.get('refinement_sub_step', 'N/A') + ')' if st.session_state.get('current_step') == 'refine_idea' else ''}")