import streamlit as st
from typing import TypedDict, Optional
from langgraph.graph import StateGraph, END
from together import Together
import os
import re
 
# Initialize Together client
client = Together(api_key="d03eaa038011423d1b8828a3ff2956511b2ca34b6cc88ffcbc4ad6314db31319")
 
def llm_invoke(prompt: str) -> str:
    """Call Together LLM with consistent config"""
    try:
        response = client.completions.create(
            model="mistralai/Mixtral-8x7B-Instruct-v0.1",
            prompt=prompt,
            max_tokens=3000,
            temperature=0.2,
            stop=["</s>"]
        )
        return response.choices[0].text.strip()
    except Exception as e:
        st.error(f"LLM API Error: {str(e)}")
        return "Error in LLM response"
 
# State Definition
class AgentState(TypedDict):
    domain: str
    goal: str
    follow_up_questions: Optional[str]
    follow_up_answers: Optional[str]
    agents_needed: Optional[str]
    code_generator: Optional[str]
    explanation: Optional[str]
 
# Agent Functions
def collect_domain_goal(state: AgentState) -> AgentState:
    """Simply pass through the domain and goal"""
    return state
 
def needs_follow_up(goal: str) -> bool:
    """Check if extra context is needed for files/data"""
    prompt = f"""
    Does this system goal require clarifying questions about input data, file structure, or formats?
 
    Goal: {goal}
 
    Respond only 'yes' or 'no'.
    """
    response = llm_invoke(prompt).strip().lower()
    return response == "yes"
 
def ask_follow_up_questions(state: AgentState) -> AgentState:
    """Generate follow-up questions if needed"""
    if needs_follow_up(state["goal"]):
        prompt = f"""
        Based on the domain: {state['domain']} and goal: {state['goal']},
        list 2–3 critical follow-up questions about:
        - Data/file formats (CSV, JSON, etc.)
        - Column names or structure
        - Required reference files (terms, conditions, etc.)
        - Any APIs or databases used
 
        Return in bullet-point format.
        """
        state["follow_up_questions"] = llm_invoke(prompt)
    else:
        state["follow_up_questions"] = "No questions needed"
    return state
 
def identify_agents(state: AgentState) -> AgentState:
    """Identify the agents needed for the system"""
    prompt = f"""
    Domain: {state['domain']}
    Goal: {state['goal']}
    Context: {state.get('follow_up_answers', 'No additional info')}
 
    List the agents needed in an agentic system. For each:
    - Name
    - Purpose
    - What it processes or validates
 
    Return as a clean markdown list.
    """
    state["agents_needed"] = llm_invoke(prompt)
    return state
 
def generate_code(state: AgentState) -> AgentState:
    """Generate the complete LangGraph application code"""
    prompt = f"""
    You are an expert Python engineer. Based on:
    - Domain: {state['domain']}
    - Goal: {state['goal']}
    - Follow-up Answers: {state.get('follow_up_answers', 'None')}
    - Agents Needed: {state.get('agents_needed', 'None')}
 
    Generate a full working LangGraph app using this structure:
    1. Imports: streamlit, langgraph, together, re, os, csv if needed
    2. State: TypedDict
    3. One clean agent function per node
    4. LangGraph construction with node/edge logic
    5. Together LLM call helper
    6. Streamlit UI to collect user input and run graph
    7. Output final result cleanly
 
    Match code style and flow of our health advisory example. Avoid AgentState subclasses. Ensure the system is executable.
    """
    state["code_generator"] = llm_invoke(prompt)
    return state
 
def explain_code(state: AgentState) -> AgentState:
    """Generate explanation for the generated code"""
    if state.get("code_generator"):
        prompt = f"Explain this code simply and clearly:\n\n{state['code_generator']}"
        state["explanation"] = llm_invoke(prompt)
    else:
        state["explanation"] = "No code to explain"
    return state
 
# Graph Construction
def create_graph():
    """Create and return the compiled graph"""
    graph = StateGraph(AgentState)
   
    # Add nodes
    graph.add_node("collect_domain_goal", collect_domain_goal)
    graph.add_node("ask_follow_up_questions", ask_follow_up_questions)
    graph.add_node("identify_agents", identify_agents)
    graph.add_node("generate_code", generate_code)
    graph.add_node("explain_code", explain_code)
   
    # Set entry point
    graph.set_entry_point("collect_domain_goal")
   
    # Add edges
    graph.add_edge("collect_domain_goal", "ask_follow_up_questions")
    graph.add_edge("ask_follow_up_questions", "identify_agents")
    graph.add_edge("identify_agents", "generate_code")
    graph.add_edge("generate_code", "explain_code")
    graph.add_edge("explain_code", END)
   
    return graph.compile()
 
# Initialize the app
app = create_graph()
 
# Streamlit Interface
def main():
    st.set_page_config(page_title="Agentic Code Generator", layout="centered")
    st.title("🤖 Auto-Build LangGraph App")
   
    st.markdown("Enter your **application domain** and a **clear goal** (what you want the app to do).")
   
    # Input fields
    domain = st.text_input("Domain (e.g. finance, warranty claims, medical):")
    goal = st.text_area("System Goal (be detailed, include what data is needed):", height=150)
   
    # Initialize session state
    if "state" not in st.session_state:
        st.session_state.state = {
            "domain": "",
            "goal": "",
            "follow_up_questions": None,
            "follow_up_answers": None,
            "agents_needed": None,
            "code_generator": None,
            "explanation": None
        }
   
    if "step" not in st.session_state:
        st.session_state.step = "initial"
   
    # Generate button
    if st.button("🚀 Generate Agentic App") and domain and goal:
        st.session_state.state.update({"domain": domain, "goal": goal})
        st.session_state.step = "processing"
       
        with st.spinner("Thinking like an architect..."):
            try:
                result = app.invoke(st.session_state.state)
                st.session_state.state.update(result)
               
                # Check if follow-up questions are needed
                if (st.session_state.state.get("follow_up_questions") and
                    st.session_state.state["follow_up_questions"] != "No questions needed"):
                    st.session_state.step = "follow_up"
                else:
                    st.session_state.step = "complete"
                   
            except Exception as e:
                st.error(f"Error processing request: {str(e)}")
                st.session_state.step = "error"
   
    # Handle follow-up questions
    if st.session_state.step == "follow_up":
        st.subheader("🔍 Follow-Up Questions")
        if st.session_state.state.get("follow_up_questions"):
            st.markdown(st.session_state.state["follow_up_questions"])
           
            answers = st.text_area("Your answers to the above:", height=150, key="follow_up_answers")
           
            if st.button("Continue with those answers"):
                st.session_state.state["follow_up_answers"] = answers
                st.session_state.step = "processing_final"
               
                with st.spinner("Building your LangGraph app..."):
                    try:
                        # Re-run the graph with the follow-up answers
                        final_result = app.invoke(st.session_state.state)
                        st.session_state.state.update(final_result)
                        st.session_state.step = "complete"
                    except Exception as e:
                        st.error(f"Error in final processing: {str(e)}")
                        st.session_state.step = "error"
   
    # Display results
    if st.session_state.step == "complete" and st.session_state.state.get("code_generator"):
        st.subheader("🎯 Identified Agents")
        if st.session_state.state.get("agents_needed"):
            st.markdown(st.session_state.state["agents_needed"])
       
        st.subheader("✅ Generated Code")
        st.code(st.session_state.state["code_generator"], language="python")
       
        st.subheader("🧠 Code Explanation")
        if st.session_state.state.get("explanation"):
            st.markdown(st.session_state.state["explanation"])
       
        # Download button
        if st.session_state.state.get("code_generator"):
            st.download_button(
                label="📥 Download Generated Code",
                data=st.session_state.state["code_generator"],
                file_name="generated_langgraph_app.py",
                mime="text/python"
            )
   
    # Reset button
    if st.button("🔄 Reset"):
        st.session_state.clear()
        st.experimental_rerun()
 
if __name__ == "__main__":
    main()