import os
import requests
from duckduckgo_search import DDGS
import gradio as gr

# ===============================
# CONFIG
# ===============================

BASE_URL = os.environ.get("GAIA_API_URL")
HF_USERNAME = os.environ.get("SPACE_AUTHOR_NAME", "jatinror")
SPACE_ID = os.environ.get("SPACE_ID", "jatinror/Final_Assignment_Template")
AGENT_CODE_URL = f"https://huggingface.co/spaces/{SPACE_ID}/tree/main"

print("Using SPACE_ID:", SPACE_ID)
print("Agent code URL:", AGENT_CODE_URL)

# ===============================
# SIMPLE SEARCH TOOL
# ===============================

def web_search(query, max_results=3):
    results = []
    with DDGS() as ddgs:
        for r in ddgs.text(query, max_results=max_results):
            results.append(r["body"])
    return "\n".join(results)

# ===============================
# DOWNLOAD FILE IF TASK HAS ONE
# ===============================

def download_task_file(task_id):
    url = f"{BASE_URL}/files/{task_id}"
    response = requests.get(url)
    if response.status_code == 200:
        file_path = f"/tmp/{task_id}"
        with open(file_path, "wb") as f:
            f.write(response.content)
        return file_path
    return None

# ===============================
# BASIC REASONING
# ===============================

def solve_question(question, task_id):
    file_path = download_task_file(task_id)
    context = ""
    if file_path and os.path.exists(file_path):
        try:
            with open(file_path, "r", errors="ignore") as f:
                context = f.read()
        except:
            context = ""
    else:
        context = web_search(question)
    return extract_answer(context)

# ===============================
# ANSWER EXTRACTION
# ===============================

def extract_answer(text):
    import re
    numbers = re.findall(r"\b\d+(?:\.\d+)?\b", text)
    if numbers:
        return numbers[0]
    words = text.split()
    return " ".join(words[:6]).strip()

# ===============================
# FETCH QUESTIONS
# ===============================

def get_questions():
    response = requests.get(f"{BASE_URL}/questions")
    response.raise_for_status()
    return response.json()

# ===============================
# SUBMIT ANSWERS
# ===============================

def submit_answers(answers):
    payload = {
        "username": HF_USERNAME,
        "agent_code": AGENT_CODE_URL,
        "answers": answers
    }
    print("Submitting payload...")
    response = requests.post(f"{BASE_URL}/submit", json=payload)
    print("Server response:", response.text)

# ===============================
# MAIN PIPELINE
# ===============================

def run_agent():
    print("Fetching GAIA questions...")
    questions = get_questions()
    answers = []
    for q in questions:
        task_id = q["task_id"]
        question = q["question"]
        print("Solving:", task_id)
        try:
            result = solve_question(question, task_id)
        except Exception as e:
            print("Error:", e)
            result = ""
        answers.append({
            "task_id": task_id,
            "submitted_answer": result.strip()
        })
    submit_answers(answers)
    print("Finished submission.")

# ===============================
# GRADIO BLOCKS UI
# ===============================

def run_pipeline():
    try:
        run_agent()
        return "✅ GAIA submission completed. Check leaderboard."
    except Exception as e:
        return f"❌ Error occurred: {str(e)}"

with gr.Blocks() as demo:
    run_button = gr.Button("Run GAIA Agent")
    output_text = gr.Textbox(label="Output", lines=4)
    
    # Link button click to function
    run_button.click(fn=run_pipeline, inputs=[], outputs=output_text)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)