import os
import requests
import io
from PIL import Image
from langchain import PromptTemplate, LLMChain
from PIL import Image, ImageDraw, ImageFont, ImageFilter
from langchain.llms import OpenAI
import openai
from g4f import Provider, Model
from langchain_g4f import G4FLLM

def set_openai_api_key(api_key):
    openai.api_key = api_key
    os.environ["OPENAI_API_KEY"] = openai.api_key

template = """Write a very short and simple children storybook suitable for children of {age} years and below.  


Your answer should be structred like this with <Text> and <Image> tags.
<Text> first page text here </Text>
<Image> descibe image here without including names so that prompt can be used to generate an image using a ML model.</Image>


NUMBER OF PAGES THE BOOK MUST HAVE:
{number_of_pages}
=============
CHARACTERS BOOK MUST HAVE:
{character_names}
Answer:"""

prompt = PromptTemplate(template=template, input_variables=["number_of_pages", "character_names", "age"])

def query(payload):
    API_URL = "https://api-inference.huggingface.co/models/prompthero/openjourney"
    headers = {"Authorization": "Bearer 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"}
    response = requests.post(API_URL, headers=headers, json=payload)
    return response.content

def query_alt(payload):
    API_URL = "https://api-inference.huggingface.co/models/stablediffusionapi/anything-v5"
    headers = {"Authorization": "Bearer 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"}
    response = requests.post(API_URL, headers=headers, json=payload)
    return response.content

def generate_plot(number_of_pages, character_names, age, selected_style, provider, selected_provider=None):
    if provider == "OpenAI":
        llm = OpenAI(temperature=0)
    elif provider == "G4F":
        if selected_provider == "Ails":
            llm = G4FLLM(
                model=Model.gpt_35_turbo,
                provider=Provider.Ails,
            )
        elif selected_provider == "You":
            llm = G4FLLM(
                model=Model.gpt_35_turbo,
                provider=Provider.You,
            )
        elif selected_provider == "GetGpt":
            llm = G4FLLM(
                model=Model.gpt_35_turbo,
                provider=Provider.GetGpt,
            )
        elif selected_provider == "DeepAi":
            llm = G4FLLM(
                model=Model.gpt_35_turbo,
                provider=Provider.DeepAi,
            )
        elif selected_provider == "Forefront":
            llm = G4FLLM(
                model=Model.gpt_35_turbo,
                provider=Provider.Forefront,
            )
        elif selected_provider == "Aichat":
            llm = G4FLLM(
                model=Model.gpt_35_turbo,
                provider=Provider.Aichat,
            )
        elif selected_provider == "Bard":
            llm = G4FLLM(
                model=Model.gpt_35_turbo,
                provider=Provider.Bard,
            )
        # Add other providers here
        else:
            raise ValueError("Invalid G4F provider selected.")
    else:
        raise ValueError("Invalid provider selected.")

    llm_chain = LLMChain(prompt=prompt, llm=llm)
    response = llm_chain.run(number_of_pages=number_of_pages, character_names=character_names, age=age)
    pages = response.split("<Text>")
    plot_result = []

    additional_texts = {    
        "Style 1" : " clean cel shaded vector art. shutterstock. behance hd by lois van baarle, artgerm, helen huang, by makoto shinkai and ilya kuvshinov, rossdraws, illustration, art by ilya kuvshinov and gustav klimt 4.",
        "Style 2" : " storybook illustration, monochromatic, white background, pinterest, trending on artstation  behance, digital illustration, vector illustration, sharp focus, pastel palette,  graphic novel, 8k",
        "Style 3" : " storybook illustration, monochromatic, white background, vector art  illustration, award winning, stunning, trending on artstation, detailed, high quality resolution, cartoon character design, adobe illustrator."
    }

    for i, page in enumerate(pages[1:]):
        text, img_text = page.split("<Image>", 1)
        selected_additional_text = additional_texts.get(selected_style, "")
        
        new_img_text = img_text.strip() + " " + selected_additional_text
        text = text.replace("</Text>", "")
        
        new_img_text = new_img_text.replace("</Image>", "")
        
        plot_result.append((i + 1, "<Text>" + text.strip() + "</Text>", "<Image>" + new_img_text + "</Image>"))
    return plot_result



def generate_storybook(plot_result):
    storybook = []
    for page_number, text, image_text in plot_result:
        image_bytes = query({"inputs": image_text})
        image = Image.open(io.BytesIO(image_bytes))
        blurred_image = image.filter(ImageFilter.GaussianBlur(8))

        draw = ImageDraw.Draw(blurred_image)

        font_size = 30
        font = ImageFont.truetype("DeliusSwashCaps-Regular.ttf", font_size)

        first_letter_font_size = 60
        first_letter_font = ImageFont.truetype("DeliusSwashCaps-Regular.ttf", first_letter_font_size)

        text_x, text_y = 50, 50
        max_width = blurred_image.width - text_x * 2

        text = text.replace("<Text>", "").replace("</Text>", "")
        wrapped_text = ""
        words = text.split()
        for i, word in enumerate(words):
            if i == 0:  
                first_letter_width = draw.textsize(word[0], font=first_letter_font)[0]
                draw.text((text_x, text_y), word[0], fill="white", font=first_letter_font, stroke_width=2, stroke_fill="black")
                text_x += first_letter_width + 10
                wrapped_text += word[1:] + " "  
            elif draw.textsize(wrapped_text + word, font=font)[0] < max_width:
                wrapped_text += word + " "
            else:
                draw.text((text_x, text_y), wrapped_text.strip(), fill="white", font=font, stroke_width=2, stroke_fill="black")
                text_y += font.getsize(wrapped_text)[1] + 10
                wrapped_text = word + " "

        draw.text((text_x, text_y), wrapped_text.strip(), fill="white", font=font, stroke_width=2, stroke_fill="black")
        combined_image = Image.new('RGB', (image.width * 2, image.height))
        combined_image.paste(blurred_image, (0, 0))
        combined_image.paste(image, (image.width, 0))

        storybook.append((page_number, combined_image))

    return storybook


def generate_book_cover(title, author, image_text):
    book_cover = []

    image_bytes = query({"inputs": image_text})
    image = Image.open(io.BytesIO(image_bytes))

    cover_image = Image.new('RGB', (image.width * 2, image.height), color='white')
    cover_image.paste(image, (0, 0))

    draw = ImageDraw.Draw(cover_image)
    font_size = 50
    title_font = ImageFont.truetype("DeliusSwashCaps-Regular.ttf", font_size)
    author_font = ImageFont.truetype("DeliusSwashCaps-Regular.ttf", 30)

    title_x, title_y = image.width + 50, 50
    author_x, author_y = image.width + 50, title_y + 100

    draw.text((title_x, title_y), title, fill="black", font=title_font)
    draw.text((author_x, author_y), "By " + author, fill="black", font=author_font)

    book_cover.append(cover_image)

    return book_cover