import streamlit as st
import os
from langchain.agents import initialize_agent, Tool
from langchain.agents.agent_types import AgentType
from langchain_community.llms.groq import ChatGroq
from dotenv import load_dotenv

# Load .env
load_dotenv()

# --- Groq API Key ---
GROQ_API_KEY = os.getenv("GROQ_API_KEY") or st.secrets.get("GROQ_API_KEY", "")
if not GROQ_API_KEY:
    st.error("❌ Please set your GROQ_API_KEY in Hugging Face Secrets.")
    st.stop()

# --- LLM Setup ---
llm = ChatGroq(api_key=GROQ_API_KEY, model_name="mixtral-8x7b-32768")

# --- Define a goal breakdown function (not lambda) ---
def goal_planner_tool(user_goal: str) -> str:
    return f"4-Week Goal: {user_goal}\n\n" \
           f"Week 1: Research and outline\n" \
           f"Week 2: Start execution with small tasks\n" \
           f"Week 3: Iterate, improve, and scale\n" \
           f"Week 4: Final review, wrap-up, and reflect\n\n" \
           f"Tip: Stay consistent and track your progress daily."

# --- Tools ---
tools = [
    Tool(
        name="GoalPlanner",
        func=goal_planner_tool,
        description="Breaks a 4-week goal into weekly action steps."
    )
]

# --- Initialize Agent ---
agent = initialize_agent(
    tools,
    llm,
    agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
    verbose=False
)

# --- Streamlit App UI ---
st.set_page_config(page_title="Goal Planner Buddy", layout="centered")

# --- Custom dark mode styling ---
custom_css = """
<style>
    .stApp {
        background-color: #121212;
        color: white;
    }
    input, textarea {
        background-color: #1e1e1e !important;
        color: white !important;
    }
    .stButton > button {
        background-color: #4CAF50;
        color: white;
    }
</style>
"""
st.markdown(custom_css, unsafe_allow_html=True)

# --- UI ---
st.title("🎯 Goal Planner Buddy")
st.subheader("Break your goals into weekly plans with AI")

goal_input = st.text_input("Enter your 4-week goal:")
if st.button("Generate Plan"):
    if goal_input:
        try:
            prompt = f"Break down this 4-week goal into weekly tasks and motivation tips: {goal_input}"
            response = agent.run(prompt)
            st.success("✅ Plan Generated:")
            st.markdown(f"```\n{response}\n```")
        except Exception as e:
            st.error(f"⚠️ Agent error: {e}")
    else:
        st.warning("Please enter a goal.")