import os
import gradio as gr
from langchain_huggingface import HuggingFaceEndpoint, ChatHuggingFace
from langchain_core.messages import HumanMessage, SystemMessage

# Load token from Hugging Face Space Secrets
hf_token = os.getenv("HF_TOKEN")
os.environ["HUGGINGFACEHUB_API_TOKEN"] = str(hf_token)

# 1. Initialize Model
repo_id = "mistralai/Mistral-7B-Instruct-v0.2" 

raw_llm = HuggingFaceEndpoint(
    repo_id=repo_id,
    task="conversational", 
    max_new_tokens=1024,
    temperature=0.7,
    huggingfacehub_api_token=hf_token
)
chat_model = ChatHuggingFace(llm=raw_llm)

# 2. PM Logic Function
def refine_prd(product_idea, focus_area):
    if not product_idea.strip():
        return "### ⚠️ Please enter a product idea to begin."

    personas = {
        "Engineering": "Focus on technical constraints, scalability, and edge cases.",
        "UX Design": "Focus on user friction, accessibility, and the 'Aha!' moment.",
        "Business/Executive": "Focus on ROI, market differentiation, and North Star metrics.",
        "Full 360° Review": "Provide a condensed analysis covering Engineering, UX, and Business risks in one view."
    }
    
    selected_persona = personas.get(focus_area)
    
    messages = [
        SystemMessage(content=f"You are a Senior PM. Review this idea from the {focus_area} lens. Goal: {selected_persona}"),
        HumanMessage(content=f"Provide: 1. Strategic Critique, 2. Three User Stories, 3. One major risk for this idea: {product_idea}")
    ]

    try:
        response = chat_model.invoke(messages)
        return response.content
    except Exception as e:
        return f"### ❌ Error\n{str(e)}"

# 3. Gradio Interface Layout
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🚀 Multi-Stakeholder PRD Refiner")
    gr.Markdown("**Product Description:** This tool acts as an automated 'Three-Amigos' review session, instantly stress-testing raw product ideas through Engineering, UX, and Business lenses to identify risks and refine requirements before development begins.")
    
    with gr.Row():
        with gr.Column(scale=1):
            idea = gr.Textbox(
                label="Raw Product Idea", 
                placeholder="Describe your feature or product idea here...",
                lines=10
            )
            persona = gr.Dropdown(
                choices=["Engineering", "UX Design", "Business/Executive", "Full 360° Review"], 
                label="Select Perspective", 
                value="Full 360° Review"
            )
            btn = gr.Button("Run Stakeholder Analysis", variant="primary")
        
        with gr.Column(scale=2):
            output = gr.Markdown(label="Refined Output")

    # 4. Examples Section
    gr.Examples(
        examples=[
            ["A mobile app that uses AI to identify bird calls and log them for citizen science.", "Engineering"],
            ["A 'Buy Now Pay Later' feature for an existing e-commerce platform.", "Business/Executive"],
            ["A smart fridge interface that suggests recipes based on expiring ingredients.", "UX Design"]
        ],
        inputs=[idea, persona]
    )
    
    btn.click(fn=refine_prd, inputs=[idea, persona], outputs=output)

if __name__ == "__main__":
    # Note: Spaces handles the port and host automatically
    demo.launch()