import json
import os
import tempfile

import wandb


def save_results(all_results: list[dict], base_filename: str = "optimization_results"):
    """
    Save optimization results to a JSON file and WandB artifact.

    Args:
        all_results: Results to save
        base_filename: Base name for the JSON file (without extension)
    """
    # Save results to a JSON file with unique name
    filename = f"{base_filename}.json"
    n = 1
    while os.path.exists(filename):
        filename = f"{base_filename}_{n:02d}.json"
        n += 1

    with open(filename, "w") as f:
        json.dump(all_results, f, indent=2)

    # Also save to wandb artifact
    results_artifact = wandb.Artifact(
        "optimization_results",
        type="results",
        description="Results from simple stories optimization",
    )
    with tempfile.NamedTemporaryFile(mode="w", suffix=".json") as tmp:
        json.dump(all_results, tmp)
        tmp.flush()
        results_artifact.add_file(tmp.name, "optimization_results.json")
    wandb.log_artifact(results_artifact)

    print("Results saved to optimization_results.json and logged to WandB")
