import os
import re
from typing import Optional, Union
from llms import build_prompt, get_completion, PROMPT_TEMPLATE, SYSTEM_PROMPT


def generate_code(problem_description: str, prompt_template: str, system_prompt: str, model: str, temperature: float):
    prompt = build_prompt(problem_description, prompt_template)
    completion = get_completion(prompt, system_prompt, model, temperature)
    return completion


def load_txt_file(file_path: str):
    with open(file_path, "r") as file:
        return file.read()

def load_problem_description(day: int):
    day_str = f"{day:02d}"
    file_path = f"day{day_str}/problem.txt"
    return load_txt_file(file_path)


def extract_code(completion: str) -> str:
    """Extracts the code from the completion, should be contained in ```python ...``` code block."""
    code_block = re.search(r"```python\s*([\s\S]*?)\s*```", completion)
    if code_block:
        return code_block.group(1)

def get_solution_file_path(model: str, day: Optional[Union[str, int]] = None):
    """Returns the path to the solution file for the given day and model. If day is None, returns the path to the solution file only."""

    if day is None:
        return f"solution_{model}.py"

    # We want it formatted properly, so we convert to int and back if its already a string for convenience
    if isinstance(day, str):
        day = int(day)

    day_str = f"{day:02d}"

    return f"day{day_str}/solution_{model}.py"


def save_code(day: int, code: str, model: str):
    """Saves the code to the solution file for the given day and model."""
    file_path = get_solution_file_path(day, model)
    with open(file_path, "w") as file:
        file.write(code)
    print(f"Saved code to {file_path}")


all_models = {
    "openai": ["gpt-4o"],
    "gemini": ["gemini-1.5-pro"],
    "anthropic": ["claude-3-5-sonnet-20241022"],
}


if __name__ == "__main__":

    for day in range(1, 26):
        problem_description = load_problem_description(day)
        print(f"***Generating code for day {day}***")

        for provider in all_models:
            for model in all_models[provider]:
                print("-" * 80)
                print(f"Generating code for {provider} {model}")

                if os.path.exists(get_solution_file_path(model=model, day=day)):
                    print(f"Skipping {provider} {model} for day {day} because it already exists")
                    continue

                prompt = build_prompt(problem_description, PROMPT_TEMPLATE)
                completion = get_completion(provider=provider, user_prompt=prompt, system_prompt=SYSTEM_PROMPT, model=model, temperature=0)
                code = extract_code(completion)
                save_code(day, code, model)
                print("-" * 80)

        print("*" * 80)