import json
import warnings
from pathlib import Path
from typing import Union

import openai


class OpenAITrainer:
    def __init__(self, model: str = "ada") -> None:
        assert model in [
            "ada",
            "babbage",
            "curie",
            "davinci",
        ], "The model must be one of the following: ada, babbage, curie, davinci."
        self.model = model

    def train(
        self,
        path_or_buf: Union[str, Path, list],
        epochs: int = 10,
        batch_size: int = 32,
    ) -> str:
        if isinstance(path_or_buf, list):
            file_path = self._prepare_training_data(path_or_buf)
        elif isinstance(path_or_buf, Path):
            file_path = path_or_buf.as_posix()
        else:
            file_path = path_or_buf

        upload_response = openai.File.create(
            file=open(file_path, "rb"),
            purpose="fine-tune",
        )
        file_id = upload_response.id
        fine_tune_response = openai.FineTune.create(
            training_file=file_id,
            model=self.model,
            n_epochs=epochs,
            batch_size=batch_size,
        )
        self.fine_tune_id = fine_tune_response.id
        warnings.warn(
            "Since the OpenAI API may take from minutes to hours depending on the size"
            " of the training data, then from now on, you'll be able to check its"
            " progress via the following command: `openai api fine_tunes.follow -i"
            f" {self.fine_tune_id}`. Once the training is completed, then you'll be"
            " able to use `OpenAIPredict` with the either the fine tune id returned,"
            " or from the model name generated by OpenAI linked to your account.",
            stacklevel=2,
        )
        return self.fine_tune_id

    def track_training(self) -> str:
        if not self.fine_tune_id:
            raise ValueError("You must first train the model.")
        return openai.FineTune.stream_events(self.fine_tune_id)

    def _prepare_training_data(self, buf: list) -> str:
        file_path = Path.cwd() / "training_data.jsonl"
        with open(file_path, "w") as f:
            for entry in buf:
                json.dump(entry, f)
                f.write("\n")
        return file_path.as_posix()
