import pathlib
import tempfile
from typing import BinaryIO, Literal
import json

import gradio as gr
from datasets import load_dataset, Dataset
from huggingface_hub import upload_file
from evaluation import evaluate_problem
from datetime import datetime


PROBLEM_TYPES = ["geometrical", "simple_to_build", "mhd_stable"]
repo_id = "cgeorgiaw/constellaration-submissions"

def submit_boundary(
    problem_type: Literal["geometrical", "simple_to_build", "mhd_stable"],
    boundary_file: BinaryIO,
) -> str:
    
    file_path = boundary_file.name
    if not file_path:
        return "Error: Uploaded file object does not have a valid file path."
    
    path_obj = pathlib.Path(file_path)
    timestamp = datetime.utcnow().isoformat()

    with (
        path_obj.open("rb") as f_in,
        tempfile.NamedTemporaryFile(delete=False, suffix=".json") as tmp,
    ):
        file_content = f_in.read()
        tmp.write(file_content)
        tmp_path = pathlib.Path(tmp.name)

        # write to dataset
        filename = f"data/{timestamp.replace(':', '-')}_{problem_type}.json"
        record = {
            "submission_time": timestamp,
            "problem_type": problem_type,
            "boundary_json": file_content.decode("utf-8"),  # Or store file path or URL
        }
        json.dump(record, tmp, indent=2)
        tmp.flush()
        upload_file(
            path_or_fileobj=f_in, 
            path_in_repo=filename,
            repo_id=repo_id,
            repo_type="dataset",
            commit_message=f"Add submission for {problem_type} at {timestamp}"
        )

    try:
        result = evaluate_problem(problem_type, str(tmp_path))
        output = str(result)
    except Exception as e:
        output = f"Error during evaluation:\n{e}"
    finally:
        tmp_path.unlink()

    return output

def gradio_interface() -> gr.Blocks:
    with gr.Blocks() as demo:
        gr.Markdown(
            """
        # Plasma Boundary Evaluation App
        Upload your plasma boundary JSON and select the problem type to get your score.
        """
        )
        with gr.Row():
            problem_type = gr.Dropdown(
                PROBLEM_TYPES, label="Problem Type", value="geometrical"
            )
            boundary_file = gr.File(label="Boundary JSON File (.json)")

        boundary_file
        output = gr.Textbox(label="Evaluation Result", lines=10)
        submit_btn = gr.Button("Evaluate")
        submit_btn.click(
            submit_boundary,
            inputs=[problem_type, boundary_file],
            outputs=output,
        )
    return demo


if __name__ == "__main__":
    gradio_interface().launch()