import os
import json
from typing import Dict

import openai

context_prompt = """
You are a creator of illustrated books building a series of scenes for your book.
Your boss asked you to write a summary and illustrations of this text:
$TRANSCRIPTION
You have to write the summary using a maximum of 7 scenes in a JSON object following these instructions:
[{"Scene": int, "Summary": "$SUMMARY_LANGUAGE str", "Illustration": "English str"}, ...] where:
"Scene": The number of the scene.
"Summary": $SUMMARY_LANGUAGE string with a summary of the scene. It should be in $SUMMARY_LANGUAGE, and it should be less than 50 words. Readers should understand it without looking at the illustration.
"Illustration": English string with a description of an illustration for this scene in less than 20 words. It should represent the scene, but it should be understandable without the context of the scene.
Just answer with the JSON object:
"""

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 = 1240,
                 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,
                        summary_language: str) -> Dict:
        gpt_prompt = context_prompt.replace("$TRANSCRIPTION", prompt)
        gpt_prompt = gpt_prompt.replace("$SUMMARY_LANGUAGE", summary_language)
        print("gpt_prompt", gpt_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
        )
        print("GPT response", response)
        scenes = json.loads(response["choices"][0]["text"])
        print("scenes", scenes)
        if (type(scenes) == list):
            scenes = {i: d for i, d in enumerate(scenes)}
        return scenes