import time

from openai import OpenAI
from typing_extensions import override
from openai import AssistantEventHandler
from pydantic import BaseModel
import json

api_key = "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"

class Criteria(BaseModel):
    explanation: str
    student_criteria_score: int
    criteria_point_total: int


class Question(BaseModel):
    explanation: str
    sub_questions: list[Criteria]
    student_question_score: int
    question_point_total: int


class GradeOutput(BaseModel):
    id: str
    graded: list[Question]
    total_student_score: int
    total_possible_points: int

class EventHandler(AssistantEventHandler):
    @override
    def on_text_created(self, text) -> None:
        print(f"\nassistant > ", end="", flush=True)

    @override
    def on_text_delta(self, delta, snapshot):
        print(delta.value, end="", flush=True)

    def on_tool_call_created(self, tool_call):
        print(f"\nassistant > {tool_call.type}\n", flush=True)

    def on_tool_call_delta(self, delta, snapshot):
        if delta.type == 'code_interpreter':
            if delta.code_interpreter.input:
                print(delta.code_interpreter.input, end="", flush=True)
            if delta.code_interpreter.outputs:
                print(f"\n\noutput >", flush=True)
                for output in delta.code_interpreter.outputs:
                    if output.type == "logs":
                        print(f"\n{output.logs}", flush=True)




class GPTGrader:
    def __init__(self, hw_path, rubric_path):
        self.hw_path = hw_path
        self.rubric_path = rubric_path
        self.client = OpenAI(api_key=api_key)
        self.context = ("You are a TA grading a college-level discrete math proof class. "
                        "Follow the rubric and grade the answers rigorously. "
                        "Grade with university-level rigor. No partial credit for any part of the rubric. "
                        "Alternate solutions are permitted as long as they hold up to our standards.")

        self.hw_string = self.read_file(self.hw_path)
        self.rubric_string = self.read_file(self.rubric_path)

    def read_file(self, file_path):
        with open(file_path, "r", encoding="ISO-8859-1") as file:
            return file.read()

    def grade(self):
        prompt = (f"Here is the rubric to use. Understand the point values for everything and "
                  f"then read the next message with the answers to grade:\n\n{self.rubric_string}\n\n"
                  f"Here is the answers to be graded:\n\n{self.hw_string}")

        response = self.client.beta.chat.completions.parse(
            model="gpt-4o",
            messages=[{"role": "system", "content": self.context},
                      {"role": "user", "content": prompt}],
            response_format=GradeOutput
        )

        json_response = json.loads(response.choices[0].message.model_dump_json())['content']
        print("LOADED:")
        event = response.choices[0].message.parsed

        return json_response