import os
import json
from typing import Dict

import openai

context_prompt = """

You are a creator of ilustraded books building a series of scenes for your book.
Your boss asked you to write a summary and the illustrations of the following text:

$TRANSCRIPTION

You have to write the summary using a maximum of 7 scenes, and using the following JSON format:

Write your answer in JSON format that has: The number of the scene, the summary for each scene, and the Illustration for each scene. 
The value for "Summary" should be in Spanish and it should not be longer than 30 words. 
The value for "Illustration" should be in English and no longer than 20 words. It should have a detail description of an illustration for this scene in English with many details, and a artistic style for the illustration that matches the text. 

Just answer with the JSON object, so your boss can easily parse it.

"""

openai.api_key = os.getenv("SECRET_KEY_OPENAI")

class TextProcessor:
    def __init__(self,
                 model: str = "text-davinci-003",
                 temperature: float = 0.7,
                 max_tokens: int = 2500,
                 top_p: int = 1,
                 frequency_penalty: int = 0,
                 presence_penalty: int = 0) -> None:
        self.model = model
        self.temperature = temperature
        self.max_tokens = max_tokens
        self.top_p = top_p
        self.frequency_penalty = frequency_penalty
        self.presence_penalty = presence_penalty
    
    def get_json_scenes(self,
                        prompt: str) -> Dict:
        gpt_prompt = context_prompt.replace("$TRANSCRIPTION", prompt)
        response = openai.Completion.create(
            model=self.model,
            prompt=gpt_prompt,
            temperature=self.temperature,
            max_tokens=self.max_tokens,
            top_p=self.top_p,
            frequency_penalty=self.frequency_penalty,
            presence_penalty=self.presence_penalty
        )
        scenes = json.loads(response["choices"][0]["text"])
        if (type(scenes) == list):
            scenes = {i: d for i, d in enumerate(scenes)}
        return scenes