import os
import requests
import json

# Load task IDs from questions_data.json
questions_file = "/home/pietro/Courses/HF_Agents/Final_Assignment_Template/questions_data.json"
with open(questions_file, "r", encoding="utf-8") as f:
    questions_data = json.load(f)

task_ids = [question["task_id"] for question in questions_data]

# Directory to save the results
output_dir = "results"
os.makedirs(output_dir, exist_ok=True)

base_url = "https://agents-course-unit4-scoring.hf.space/files/{}"
headers = {"accept": "application/json"}

for task_id in task_ids:
    url = base_url.format(task_id)
    response = requests.get(url, headers=headers)
    
    # Extract file name from Content-Disposition header if available
    content_disposition = response.headers.get("content-disposition", "")
    if "filename=" in content_disposition:
        file_name = content_disposition.split("filename=")[-1].strip('"')
    else:
        file_name = f"{task_id}.unknown"  # Default file name if not provided
    
    # Determine file extension based on Content-Type header
    content_type = response.headers.get("content-type", "")
    if "image/png" in content_type:
        file_extension = ".png"
    elif "application/json" in content_type:
        file_extension = ".json"
    elif "audio/mpeg" in content_type:
        file_extension = ".mp3"
    elif "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet" in content_type:
        file_extension = ".xlsx"
    elif "text/plain" in content_type:
        file_extension = ".txt"
    elif "application/x-python-code" in content_type or file_name.endswith(".py"):
        file_extension = ".py"
    else:
        file_extension = ""  # Unknown file type
    
    # Save the file with the appropriate name and extension
    output_path = os.path.join(output_dir, file_name if file_extension else f"{task_id}{file_extension}")
    with open(output_path, "wb") as f:
        f.write(response.content)
    print(f"Saved: {output_path}")