import os
import gradio as gr
import requests
import pandas as pd
from smolagents import CodeAgent, DuckDuckGoSearchTool

HF_TOKEN = os.getenv("HF_TOKEN")

# Auto-detect model class
try:
    from smolagents import InferenceClientModel
    ModelClass = InferenceClientModel
    print("✅ Using InferenceClientModel")
except ImportError:
    from smolagents import ApiModel
    ModelClass = ApiModel
    print("✅ Using ApiModel fallback")

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

class BasicAgent:
    def __init__(self):
        print("Initializing GAIA Agent...")
        if not HF_TOKEN or not HF_TOKEN.startswith("hf_"):
            raise ValueError("HF_TOKEN not found in Secrets!")

        # Using a SMALL & FREE model to avoid payment error
        self.model = ModelClass(
            model_id="HuggingFaceTB/SmolLM2-1.7B-Instruct",   # Very small & usually free
            token=HF_TOKEN
        )
        print(f"✅ Model loaded: {type(self.model).__name__} - SmolLM2 1.7B")

        self.agent = CodeAgent(
            tools=[DuckDuckGoSearchTool()],
            model=self.model,
            max_steps=15,          # Reduced for speed
            verbosity_level=1
        )

    def __call__(self, question: str) -> str:
        prompt = f"""
You are solving a GAIA question. Be precise.
Use search tool if needed.
Return ONLY the final answer. No explanation.

Question: {question}
"""

        try:
            answer = self.agent.run(prompt)
            answer = str(answer).strip()

            if "FINAL ANSWER:" in answer.upper():
                answer = answer.split("FINAL ANSWER:", 1)[-1].strip()

            return answer
        except Exception as e:
            print(f"Agent error: {e}")
            return f"AGENT_ERROR: {str(e)[:80]}"


def run_and_submit_all(profile: gr.OAuthProfile | None):
    if not profile:
        return "Please log in first!", None

    username = profile.username
    space_id = os.getenv("SPACE_ID")

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

    try:
        agent = BasicAgent()
    except Exception as e:
        return f"Failed to init agent: {e}", None

    agent_code = f"https://huggingface.co/spaces/{space_id}/tree/main"

    # Fetch questions
    try:
        resp = requests.get(questions_url, timeout=20)
        resp.raise_for_status()
        questions_data = resp.json()
        print(f"Fetched {len(questions_data)} questions")
    except Exception as e:
        return f"Error fetching questions: {e}", None

    # Run agent
    results_log = []
    answers_payload = []

    for item in questions_data:
        task_id = item.get("task_id")
        question = item.get("question")
        if not task_id or not question:
            continue
        try:
            ans = agent(question)
            answers_payload.append({"task_id": task_id, "submitted_answer": ans})
            results_log.append({
                "Task ID": task_id,
                "Question": question[:100] + "...",
                "Submitted Answer": str(ans)[:200]
            })
        except Exception as e:
            err = f"ERROR: {e}"
            answers_payload.append({"task_id": task_id, "submitted_answer": err})
            results_log.append({"Task ID": task_id, "Question": question[:100]+"...", "Submitted Answer": err})

    # Submit
    try:
        submission_data = {"username": username, "agent_code": agent_code, "answers": answers_payload}
        resp = requests.post(submit_url, json=submission_data, timeout=180)
        resp.raise_for_status()
        data = resp.json()
        status = f"✅ SUCCESS!\nScore: {data.get('score', 'N/A')}%\nCorrect: {data.get('correct_count', '?')}/{data.get('total_attempted', '?')}"
        return status, pd.DataFrame(results_log)
    except Exception as e:
        return f"Submission failed: {e}", pd.DataFrame(results_log)


# UI
with gr.Blocks() as demo:
    gr.Markdown("# GAIA Agent (Free Tier)")
    gr.Markdown("Using small model to avoid payment limits.")

    gr.LoginButton()
    btn = gr.Button("🚀 Run Evaluation & Submit", variant="primary", size="large")

    status = gr.Textbox(label="Status / Score", lines=10)
    table = gr.DataFrame(label="Results")

    btn.click(run_and_submit_all, outputs=[status, table])


if __name__ == "__main__":
    print("=== GAIA Agent Starting (Small Model) ===")
    demo.launch(debug=True, share=False)