import os
import time
import gradio as gr
import requests
import pandas as pd
from dotenv import load_dotenv
from agent import LangGraphGAIAAgent

# Satisfy Hugging Face ZeroGPU startup scanner if running on a ZeroGPU hardware tier
try:
    import spaces
    @spaces.GPU(duration=1)
    def _zerogpu_startup():
        return None
except Exception:
    pass

load_dotenv()

# --- Constants ---
DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"

def run_single_question(question: str, task_id: str = "", file_name: str = "", groq_api_key: str = "", model_name: str = "openai/gpt-oss-20b") -> str:
    """Run the agent on a single question for rapid testing and debugging."""
    if not question.strip():
        return "Please enter a question or select one from the list."
    
    key = groq_api_key.strip() or os.getenv("GROQ_API_KEY", "")
    if not key:
        return "Error: GROQ_API_KEY is required. Please provide it in the API settings box."
    
    selected_model = model_name.split(" ")[0].strip() if model_name else "openai/gpt-oss-120b"
    try:
        agent = LangGraphGAIAAgent(groq_api_key=key, model_name=selected_model)
        answer = agent(question.strip(), task_id=task_id.strip(), file_name=file_name.strip())
        return answer
    except Exception as e:
        return f"Error executing agent: {e}"


def fetch_questions_for_dropdown():
    """Fetch benchmark questions to populate the test dropdown."""
    try:
        response = requests.get(f"{DEFAULT_API_URL}/questions", timeout=15)
        response.raise_for_status()
        data = response.json()
        choices = []
        for item in data:
            tid = item.get("task_id", "")
            q = item.get("question", "")
            fn = item.get("file_name", "")
            fn_str = f" [{fn}]" if fn else ""
            label = f"{q[:75]}...{fn_str} (ID: {tid[:8]})"
            choices.append((label, f"{tid}|||{fn}|||{q}"))
        return gr.update(choices=choices, value=choices[0][1] if choices else None)
    except Exception as e:
        return gr.update(choices=[(f"Failed to load questions: {e}", "")])


def select_question_from_dropdown(selected_val: str):
    """Update question, task_id, and file_name inputs when a question is chosen."""
    if not selected_val or "|||" not in selected_val:
        return "", "", ""
    parts = selected_val.split("|||")
    task_id = parts[0] if len(parts) > 0 else ""
    file_name = parts[1] if len(parts) > 1 else ""
    question = parts[2] if len(parts) > 2 else ""
    return question, task_id, file_name


def run_and_submit_all(
    profile: gr.OAuthProfile | None,
    groq_api_key: str = "",
    model_name: str = "openai/gpt-oss-20b",
    local_username: str = "",
    progress=gr.Progress()
):
    """
    Fetches all 20 questions, runs the LangGraph Agent on each, submits all answers,
    and displays the results and official score.
    """
    space_id = os.getenv("SPACE_ID")

    if profile:
        username = f"{profile.username}"
        print(f"User logged in via OAuth: {username}")
    elif local_username and local_username.strip():
        username = local_username.strip()
        print(f"Using local/manual username: {username}")
    else:
        # Fallback for local testing
        local_user = os.getenv("HF_USERNAME")
        if local_user:
            username = local_user
            print(f"Using local HF_USERNAME: {username}")
        else:
            print("User not logged in.")
            return "Please Login to Hugging Face or provide your username before submitting.", None

    key = groq_api_key.strip() or os.getenv("GROQ_API_KEY", "")
    if not key:
        return "Error: GROQ_API_KEY is required. Please provide it in the API settings box.", None

    api_url = DEFAULT_API_URL
    questions_url = f"{api_url}/questions"
    submit_url = f"{api_url}/submit"

    selected_model = model_name.split(" ")[0].strip() if model_name else "openai/gpt-oss-120b"
    # 1. Instantiate LangGraph Agent
    try:
        agent = LangGraphGAIAAgent(groq_api_key=key, model_name=selected_model)
    except Exception as e:
        print(f"Error instantiating agent: {e}")
        return f"Error initializing agent: {e}", None

    agent_code = f"https://huggingface.co/spaces/{space_id}/tree/main" if space_id else "https://huggingface.co"
    print(f"Agent code repository: {agent_code}")

    # 2. Fetch Questions
    print(f"Fetching questions from: {questions_url}")
    try:
        response = requests.get(questions_url, timeout=15)
        response.raise_for_status()
        questions_data = response.json()
        if not questions_data:
            return "Fetched questions list is empty or invalid format.", None
        print(f"Fetched {len(questions_data)} questions.")
    except Exception as e:
        return f"Error fetching questions: {e}", None

    # 3. Run Agent on all questions
    results_log = []
    answers_payload = []
    total_q = len(questions_data)
    print(f"Running LangGraph Agent on {total_q} questions using {model_name}...")

    for idx, item in enumerate(questions_data):
        task_id = item.get("task_id")
        question_text = item.get("question")
        file_name = item.get("file_name", "")

        if not task_id or question_text is None:
            continue

        progress((idx) / total_q, desc=f"[{idx+1}/{total_q}] Solving: {question_text[:35]}...")

        try:
            submitted_answer = agent(question_text, task_id=task_id, file_name=file_name)
            answers_payload.append({
                "task_id": task_id,
                "submitted_answer": submitted_answer,
                "model_answer": submitted_answer
            })
            results_log.append({
                "Task ID": task_id[:8] + "...",
                "Question": question_text,
                "Attachment": file_name or "-",
                "Submitted Answer": submitted_answer
            })
            print(f"[{idx+1}/{total_q}] Task {task_id[:8]} -> {submitted_answer}")
            time.sleep(3)  # Pacing to stay comfortably within Groq ITPM free tier limits
        except Exception as e:
            print(f"Error running agent on task {task_id}: {e}")
            results_log.append({
                "Task ID": task_id[:8] + "...",
                "Question": question_text,
                "Attachment": file_name or "-",
                "Submitted Answer": f"ERROR: {e}"
            })

    progress(1.0, desc="Submitting answers to leaderboard...")

    if not answers_payload:
        return "Agent did not produce any answers to submit.", pd.DataFrame(results_log)

    # 4. Prepare Submission
    submission_data = {
        "username": username.strip(),
        "agent_code": agent_code,
        "answers": answers_payload
    }
    print(f"Submitting {len(answers_payload)} answers for user '{username}' to {submit_url}...")

    # 5. Submit to Scoring Server
    try:
        response = requests.post(submit_url, json=submission_data, timeout=60)
        response.raise_for_status()
        result_data = response.json()
        final_status = (
            f"🎉 Submission Successful!\n"
            f"User: {result_data.get('username')}\n"
            f"Overall Score: {result_data.get('score', 'N/A')}%\n"
            f"Correct Answers: {result_data.get('correct_count', '?')}/{result_data.get('total_attempted', '?')}\n"
            f"Leaderboard Message: {result_data.get('message', 'No message received.')}"
        )
        print("Submission successful.")
        results_df = pd.DataFrame(results_log)
        return final_status, results_df
    except requests.exceptions.HTTPError as e:
        error_detail = f"Server responded with status {e.response.status_code}."
        try:
            error_json = e.response.json()
            error_detail += f" Detail: {error_json.get('detail', e.response.text)}"
        except Exception:
            error_detail += f" Response: {e.response.text[:500]}"
        return f"Submission Failed: {error_detail}", pd.DataFrame(results_log)
    except Exception as e:
        return f"An unexpected error occurred during submission: {e}", pd.DataFrame(results_log)


# --- Build Gradio Interface ---
with gr.Blocks(title="LangGraph GAIA Agent - HF Course Final Assignment") as demo:
    gr.Markdown("# 🤖 LangGraph Agent for GAIA Benchmark (Unit 4 Final Assignment)")
    gr.Markdown(
        """
        This application runs an autonomous **LangGraph ReAct Agent** powered by Groq's **`openai/gpt-oss-20b`** (with automatic fallback to `llama-3.3-70b-versatile`).
        Equipped with **DuckDuckGo Web Search**, an isolated **Python Code Interpreter**, and **Task File Inspection**, it is engineered to solve GAIA Level 1 questions and output exact-match answers.
        """
    )

    with gr.Accordion("⚙️ Groq API & Model Configuration", open=True):
        with gr.Row():
            groq_key_input = gr.Textbox(
                label="Groq API Key",
                placeholder="Enter your gsk_... key (or configure GROQ_API_KEY in Space Secrets / .env)",
                type="password",
                value=os.getenv("GROQ_API_KEY", ""),
                scale=3
            )
            model_selector = gr.Dropdown(
                label="Groq Model",
                choices=["openai/gpt-oss-120b (Recommended)", "openai/gpt-oss-20b", "qwen/qwen3.8-27b"],
                value="openai/gpt-oss-120b (Recommended)",
                scale=1
            )

    with gr.Tabs():
        with gr.TabItem("🚀 Official Evaluation & Submission"):
            gr.Markdown(
                """
                **Evaluation Instructions:**
                1. Authenticate with your Hugging Face account below (uses OAuth on HF Spaces or your HF username locally).
                2. Click **Run Evaluation & Submit All Answers**. The agent will iterate through all 20 questions, execute tools as needed, and submit answers to the official scoring API.
                3. Passing mark: **≥ 30%** (at least 6/20 exact match).
                """
            )
            is_space_env = bool(os.getenv("SPACE_ID") or os.getenv("HF_TOKEN"))
            if is_space_env:
                gr.LoginButton()
                local_username_box = gr.Textbox(visible=False, value="")
            else:
                local_username_box = gr.Textbox(
                    label="Hugging Face Username (for Leaderboard submission)",
                    placeholder="e.g. your-hf-username",
                    value=os.getenv("HF_USERNAME", "")
                )
                gr.Markdown("*(Running locally: OAuth Login Button will automatically appear when deployed to Hugging Face Spaces)*")

            run_all_button = gr.Button("▶️ Run Evaluation & Submit All Answers", variant="primary", size="lg")
            status_output = gr.Textbox(label="Submission Result / Score", lines=6, interactive=False)
            results_table = gr.DataFrame(label="Questions and Agent Answers", wrap=True)

            run_all_button.click(
                fn=run_and_submit_all,
                inputs=[groq_key_input, model_selector, local_username_box],
                outputs=[status_output, results_table]
            )

        with gr.TabItem("🧪 Test Single Question"):
            gr.Markdown("Select a question from the live benchmark or type your own question to test the LangGraph agent in real-time.")
            with gr.Row():
                fetch_q_btn = gr.Button("📥 Load Benchmark Questions", scale=1)
                preset_dropdown = gr.Dropdown(label="Select Benchmark Question", choices=[], scale=3)

            test_q_input = gr.Textbox(label="Question", lines=3, placeholder="Enter question here...")
            with gr.Row():
                test_task_id = gr.Textbox(label="Task ID (optional)", placeholder="e.g. 8e867cd7-cff9-4e6c-867a-ff5ddc2550be")
                test_file_name = gr.Textbox(label="Attached File Name (optional)", placeholder="e.g. data.xlsx or script.py")

            run_single_btn = gr.Button("🔍 Run Agent on This Question", variant="secondary")
            single_answer_output = gr.Textbox(label="Agent Sanitized Exact-Match Answer", lines=3)

            fetch_q_btn.click(fn=fetch_questions_for_dropdown, outputs=preset_dropdown)
            preset_dropdown.change(
                fn=select_question_from_dropdown,
                inputs=preset_dropdown,
                outputs=[test_q_input, test_task_id, test_file_name]
            )
            run_single_btn.click(
                fn=run_single_question,
                inputs=[test_q_input, test_task_id, test_file_name, groq_key_input, model_selector],
                outputs=single_answer_output
            )

        with gr.TabItem("ℹ️ Architecture & Tooling Details"):
            gr.Markdown(
                """
                ### Agent Architecture
                - **Orchestration**: Built on **LangGraph** using a reactive cycle (`Model Reasoning -> Tool Execution -> Re-evaluation`).
                - **Inference Engine**: Hosted on **Groq** using `openai/gpt-oss-20b` (OpenAI's open-weight MoE reasoning model) achieving ~1,000 tokens/sec for rapid multi-step reasoning.
                - **Tools**:
                  - `web_search`: DuckDuckGo real-time search for sports facts, Wikipedia promotion records, and citations.
                  - `python_interpreter`: Sandboxed execution for numerical calculation, date difference arithmetic, and Pandas dataframe querying.
                  - `download_task_file`: Fetches files associated with task IDs from the course scoring server.
                - **Exact Match Sanitizer**: Strips formatting, conversational prefixes (`Final Answer:`, `The answer is...`), bolding, and trailing punctuation to maximize GAIA exact match scoring.
                """
            )

if __name__ == "__main__":
    print("\n" + "-"*30 + " LangGraph Agent App Starting " + "-"*30)
    space_host_startup = os.getenv("SPACE_HOST")
    space_id_startup = os.getenv("SPACE_ID")

    if space_host_startup:
        print(f"✅ SPACE_HOST found: {space_host_startup}")
        url = space_host_startup if space_host_startup.startswith("http") else f"https://{space_host_startup}"
        print(f"   Runtime URL: {url}")
    else:
        print("ℹ️  SPACE_HOST not found (running locally).")

    if space_id_startup:
        print(f"✅ SPACE_ID found: {space_id_startup}")
        print(f"   Repo URL: https://huggingface.co/spaces/{space_id_startup}")
        print(f"   Repo Tree: https://huggingface.co/spaces/{space_id_startup}/tree/main")
    else:
        print("ℹ️  SPACE_ID not found (running locally).")

    print("-"*(60 + len(" LangGraph Agent App Starting ")) + "\n")
    demo.launch(ssr_mode=False)