"""
Educational AI Agent — works on Streamlit Cloud, Hugging Face Spaces, and Kaggle.
Uses Groq API for fast inference (free tier available).
"""

import os
import json
import re
import streamlit as st
from groq import Groq

# ── Page config ───────────────────────────────────────────────────────────────
st.set_page_config(
    page_title="EduAgent — AI Learning Assistant",
    page_icon="🎓",
    layout="wide",
)

# ── Minimal custom CSS ─────────────────────────────────────────────────────────
st.markdown("""
<style>
  /* Code blocks */
  .stCodeBlock { border-left: 3px solid #7c3aed; }
  /* Tool call badge */
  .tool-badge {
    display: inline-block;
    background: #ede9fe;
    color: #5b21b6;
    border-radius: 6px;
    padding: 2px 10px;
    font-size: 0.78rem;
    font-weight: 600;
    margin-bottom: 4px;
  }
  .result-box {
    background: #f5f3ff;
    border-left: 3px solid #7c3aed;
    padding: 8px 14px;
    border-radius: 0 6px 6px 0;
    font-size: 0.9rem;
    margin-bottom: 8px;
  }
</style>
""", unsafe_allow_html=True)


# ══════════════════════════════════════════════════════════════════════════════
#  1. TOOL DEFINITIONS  (sent to the LLM so it knows what it can call)
# ══════════════════════════════════════════════════════════════════════════════

TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "explain_concept",
            "description": (
                "Explain an educational concept clearly at a chosen difficulty level. "
                "Use this when the user asks 'what is', 'explain', 'define', or similar."
            ),
            "parameters": {
                "type": "object",
                "properties": {
                    "concept": {"type": "string", "description": "The concept to explain"},
                    "level": {
                        "type": "string",
                        "enum": ["beginner", "intermediate", "advanced"],
                        "description": "Target difficulty level",
                    },
                    "subject": {"type": "string", "description": "Subject area, e.g. 'math', 'physics'"},
                },
                "required": ["concept", "level"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "generate_quiz",
            "description": "Generate multiple-choice quiz questions to test understanding of a topic.",
            "parameters": {
                "type": "object",
                "properties": {
                    "topic": {"type": "string", "description": "Topic to quiz on"},
                    "num_questions": {"type": "integer", "description": "Number of questions (1-5)"},
                    "difficulty": {
                        "type": "string",
                        "enum": ["easy", "medium", "hard"],
                    },
                },
                "required": ["topic", "num_questions"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "solve_problem",
            "description": (
                "Solve a math, science, or logic problem step by step. "
                "Show all working clearly."
            ),
            "parameters": {
                "type": "object",
                "properties": {
                    "problem": {"type": "string", "description": "The problem statement"},
                    "subject": {"type": "string", "description": "Subject area"},
                },
                "required": ["problem"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "create_study_plan",
            "description": "Create a structured study plan for learning a subject or topic.",
            "parameters": {
                "type": "object",
                "properties": {
                    "subject": {"type": "string", "description": "Subject or topic to study"},
                    "duration_weeks": {"type": "integer", "description": "Available weeks"},
                    "goal": {"type": "string", "description": "Learning goal or exam target"},
                },
                "required": ["subject", "duration_weeks"],
            },
        },
    },
]


# ══════════════════════════════════════════════════════════════════════════════
#  2. TOOL EXECUTION  (local Python functions the agent can call)
# ══════════════════════════════════════════════════════════════════════════════

def execute_tool(tool_name: str, tool_args: dict, client: Groq, model: str) -> str:
    """Route tool calls to local implementations or a second LLM call."""

    # For educational tools, we call the LLM again with a focused prompt.
    # In a real app you could swap these for database lookups, calculators, etc.

    prompts = {
        "generate_quiz": (
            f"Create {tool_args.get('num_questions', 3)} multiple-choice questions about "
            f"'{tool_args.get('topic', 'the requested topic')}' at {tool_args.get('difficulty','medium')} difficulty. "
            "Format each question as:\n**Q1.** ...\na) ...\nb) ...\nc) ...\nd) ...\n✅ Answer: ..."
        ),
        "solve_problem": (
            f"Solve this step-by-step:\n{tool_args.get('problem', tool_args)}\n"
            "Show every step clearly. Box the final answer."
        ),
        "create_study_plan": (
            f"Create a {tool_args.get('duration_weeks', 4)}-week study plan for '{tool_args.get('subject', 'the subject')}'. "
            f"Goal: {tool_args.get('goal','general mastery')}. "
            "Use a weekly breakdown with topics and practice tasks."
        ),
        "explain_concept": (
            f"Explain '{tool_args['concept']}' at {tool_args.get('level','intermediate')} level"
            f"{' in ' + tool_args['subject'] if tool_args.get('subject') else ''}. "
            "Use clear language, an analogy, and a short example. Markdown ok."
        ),
    }

    prompt = prompts.get(tool_name, f"Handle this: {tool_args}")

    response = client.chat.completions.create(
        model=model,
        messages=[{"role": "user", "content": prompt}],
        max_tokens=1024,
        temperature=0.7,
    )
    return response.choices[0].message.content


# ══════════════════════════════════════════════════════════════════════════════
#  3. AGENT LOOP  (single turn: may call multiple tools before final answer)
# ══════════════════════════════════════════════════════════════════════════════

def run_agent(user_message: str, history: list, client: Groq, model: str):
    """
    Agentic loop:
    1. Send user message + history to LLM with tools.
    2. If LLM wants to call tools → execute them → feed results back.
    3. Repeat until LLM returns a final text answer.

    Yields (type, content) tuples for streaming UI updates.
    """
    messages = [
        {
            "role": "system",
            "content": (
                "You are EduAgent, a friendly and expert educational AI tutor. "
                "You have tools to explain concepts, generate quizzes, solve problems, "
                "and create study plans. Use the most appropriate tool for each request. "
                "You may chain multiple tools if needed. Always be encouraging and clear."
            ),
        }
    ] + history + [{"role": "user", "content": user_message}]

    try:
        response = client.chat.completions.create(
            model=model,
            messages=messages,
            tools=TOOLS,
            tool_choice="auto",
            max_tokens=2048,
            temperature=0.7,
        )
    except BadRequestError as e:
        if "decommissioned" in str(e):
            yield ("answer", "⚠️ This model has been retired by Groq. Please select a different model from the sidebar.", [])
            return
        raise

    tool_calls_log = []

    for _iteration in range(5):  # max tool-call depth guard
        response = client.chat.completions.create(
            model=model,
            messages=messages,
            tools=TOOLS,
            tool_choice="auto",
            max_tokens=2048,
            temperature=0.7,
        )

        msg = response.choices[0].message

        # ── No tool calls → final answer ──────────────────────────────────────
        if not msg.tool_calls:
            yield ("answer", msg.content, tool_calls_log)
            return

        # ── Process each tool call ────────────────────────────────────────────
        messages.append({"role": "assistant", "content": msg.content, "tool_calls": msg.tool_calls})

        for tc in msg.tool_calls:
            fn_name = tc.function.name
            fn_args = json.loads(tc.function.arguments)

            yield ("tool_start", fn_name, fn_args)

            result = execute_tool(fn_name, fn_args, client, model)

            yield ("tool_result", fn_name, result)
            tool_calls_log.append({"tool": fn_name, "args": fn_args, "result": result})

            messages.append({
                "role": "tool",
                "tool_call_id": tc.id,
                "content": result,
            })

    yield ("answer", "I reached the maximum reasoning steps. Please rephrase your question.", tool_calls_log)


# ══════════════════════════════════════════════════════════════════════════════
#  4. STREAMLIT UI
# ══════════════════════════════════════════════════════════════════════════════

def get_api_key() -> str | None:
    """Retrieve Groq API key from env or Streamlit secrets (all platforms)."""
    # Streamlit Cloud / HF Spaces: set in app secrets
    try:
        return st.secrets["GROQ_API_KEY"]
    except Exception:
        pass
    # Kaggle / local: environment variable
    return os.environ.get("GROQ_API_KEY")


def main():
    # ── Sidebar ────────────────────────────────────────────────────────────────
    with st.sidebar:
        st.title("⚙️ Settings")

        api_key_input = st.text_input(
            "Groq API Key",
            type="password",
            placeholder="gsk_...",
            help="Get a free key at console.groq.com",
        )
        api_key = api_key_input or get_api_key()

        model = st.selectbox(
            "Model",
            [
               "llama-3.3-70b-versatile",        # best quality, 1K req/day free
               "llama-3.1-8b-instant",            # fastest, 14.4K req/day free ← safest for testing
               "gemma2-9b-it",                    # Google Gemma, 15K TPM free
               "meta-llama/llama-4-scout-17b-16e-instruct",
            ],
            index=1,
            help="All models run on Groq's ultra-fast inference.",
        )

        st.divider()
        st.markdown("**Try asking:**")
        examples = [
            "Explain neural networks for a beginner",
            "Quiz me on photosynthesis (3 questions)",
            "Solve: integrate x² from 0 to 3",
            "Make a 4-week Python study plan",
            "Explain recursion then quiz me on it",
        ]
        for ex in examples:
            if st.button(ex, use_container_width=True, key=ex):
                st.session_state["pending_input"] = ex

        st.divider()
        if st.button("🗑️ Clear chat", use_container_width=True):
            st.session_state["messages"] = []
            st.session_state["history"] = []
            st.rerun()

        st.caption("Powered by [Groq](https://groq.com) · Free tier available")

    # ── Main area ──────────────────────────────────────────────────────────────
    st.title("🎓 EduAgent")
    st.caption("An AI learning assistant with tool-use — explains, quizzes, solves, and plans.")

    if not api_key:
        st.warning(
            "Add your **Groq API key** in the sidebar to start. "
            "Get one free at [console.groq.com](https://console.groq.com).",
            icon="🔑",
        )
        st.stop()

    # Init state
    if "messages" not in st.session_state:
        st.session_state["messages"] = []   # rendered chat
    if "history" not in st.session_state:
        st.session_state["history"] = []    # LLM-format history

    # Render existing chat
    for msg in st.session_state["messages"]:
        with st.chat_message(msg["role"]):
            st.markdown(msg["content"])

    # Handle sidebar example buttons
    pending = st.session_state.pop("pending_input", None)

    user_input = st.chat_input("Ask anything educational…") or pending
    if not user_input:
        st.stop()

    # Show user message
    with st.chat_message("user"):
        st.markdown(user_input)
    st.session_state["messages"].append({"role": "user", "content": user_input})

    # Run agent
    client = Groq(api_key=api_key)
    final_answer = ""
    tool_summaries = []

    with st.chat_message("assistant"):
        status_placeholder = st.empty()
        answer_placeholder = st.empty()

        for event in run_agent(user_input, st.session_state["history"], client, model):
            etype = event[0]

            if etype == "tool_start":
                _, tool_name, tool_args = event
                readable = tool_name.replace("_", " ").title()
                status_placeholder.markdown(
                    f'<span class="tool-badge">🔧 {readable}</span>',
                    unsafe_allow_html=True,
                )

            elif etype == "tool_result":
                _, tool_name, result = event
                readable = tool_name.replace("_", " ").title()
                tool_summaries.append(f"**{readable}** → result ready ✓")

            elif etype == "answer":
                _, final_answer, _ = event
                status_placeholder.empty()
                answer_placeholder.markdown(final_answer)

        # Show tool trace in expander if tools were used
        if tool_summaries:
            with st.expander("🔍 Agent reasoning trace", expanded=False):
                for s in tool_summaries:
                    st.markdown(s)

    # Update session history (for multi-turn)
    st.session_state["messages"].append({"role": "assistant", "content": final_answer})
    st.session_state["history"].append({"role": "user", "content": user_input})
    st.session_state["history"].append({"role": "assistant", "content": final_answer})

    # Keep history bounded (last 10 turns = 20 messages)
    if len(st.session_state["history"]) > 20:
        st.session_state["history"] = st.session_state["history"][-20:]


if __name__ == "__main__":
    main()