import os
import json
from bs4 import BeautifulSoup

def extract_test_cases(html_file):
    """Extracts input/output test cases from a problem HTML file."""
    with open(html_file, "r", encoding="utf-8") as f:
        soup = BeautifulSoup(f, "html.parser")
    
    # Find all <pre> blocks (sample test cases)
    pre_blocks = [pre.get_text().strip() for pre in soup.find_all("pre")]

    # Assume inputs and outputs alternate
    inputs, outputs = pre_blocks[::2], pre_blocks[1::2]

    return {"input": inputs, "output": outputs}

def process_problem_descriptions(folder_path):
    """Processes all HTML files in a folder and generates JSON."""
    problems = {}

    for filename in os.listdir(folder_path):
        if filename.endswith(".html"):
            problem_id = filename.replace(".html", "")  # Extract problem ID
            file_path = os.path.join(folder_path, filename)
            
            test_cases = extract_test_cases(file_path)
            problems[problem_id] = {"public_tests": test_cases}

    return problems

# Example usage
FOLDER_PATH = "path/to/your/html/files"  # Change this to the actual folder
problems_json = process_problem_descriptions(FOLDER_PATH)

# Save to JSON file
with open("problems.json", "w", encoding="utf-8") as f:
    json.dump(problems_json, f, indent=4)

print("Test cases extracted and saved to problems.json!")