from fpdf import FPDF
import gradio as gr
from urllib.parse import quote
import os
import openai
from gradio_client import Client
import requests
from PIL import Image
from io import BytesIO


def generate_image_url(keywords):
    truncated_keywords = keywords[:20]
    return f"https://image.pollinations.ai/prompt/{quote(truncated_keywords)}"

def download_image(image_url):
    response = requests.get(image_url)
    img = Image.open(BytesIO(response.content))
    local_image_path = f"/tmp/{image_url.split('/')[-1]}.png"
    img.save(local_image_path)
    return local_image_path
    
class PDF(FPDF):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.add_font('malgun', '', '/home/user/app/malgun.ttf', uni=True)
        self.set_font('/home/user/app/malgun.ttf', '', 12)
        
    def chapter_title(self, num, label):
        self.set_font('/home/user/app/malgun.ttf', '', 15)
        self.cell(0, 10, f"Chapter {num} : {label}", 0, 1, 'C')
        self.ln(10)

    def chapter_body(self, body, image_url):
        self.set_font('/home/user/app/malgun.ttf', '', 12)
        self.multi_cell(0, 10, body)
        self.image(image_url, x=10, w=100)
        self.ln(60)

    def add_chapter(self, num, title, body, image_url):
        local_image_path = download_image(image_url)
        self.add_page()
        self.chapter_title(num, title)
        self.chapter_body(body, local_image_path)

def save_as_pdf():
    global chapters
    pdf_path = "/mnt/data/chapters.pdf"
    pdf = PDF()
    for i, chapter in enumerate(chapters, start=1):
        pdf.add_chapter(i, f"Chapter {i}", chapter["story"], chapter["image"])
    pdf.output(pdf_path)
    return pdf_path

def create_markdown_table():
    global chapters
    markdown_table = "| 챕터 | 이야기 | 그림 |\n|-------|-------|------|\n"
    for i, chapter in enumerate(chapters, start=1):
        markdown_table += f"| Chapter {i} | {chapter['story']} | ![Chapter {i} Image]({chapter['image']}) |\n"
    return markdown_table

OPENAI_API_KEY = os.environ.get('OPENAI_API_KEY')
clipi_client = Client("https://fffiloni-clip-interrogator-2.hf.space/")

chapter_num = 1
story_intro = ""
current_story = ""
current_image_url = ""
chapters = []

def next_chapter(audience, keyword, protagonist):
    global chapter_num, current_story, current_image_url, chapters
    current_image_url = generate_image_url(current_story)
    gr.Info(f'Chapter {chapter_num}를 생성하고 있습니다...')
    chapter_prompt = f"{story_intro}\n\nKeyword: {keyword}\nProtagonist: {protagonist}\n\n![Chapter {chapter_num} Image]({current_image_url})\n\nChapter {chapter_num} 내용을 만들어줘."
    chat_completion = openai.ChatCompletion.create(model="gpt-3.5-turbo-16k", messages=[{"role": "user", "content": chapter_prompt}])
    current_story = chat_completion.choices[0].message.content
    chapters.append({"story": current_story, "image": current_image_url})
    chapter_num += 1
    return current_story, current_image_url, create_markdown_table()

def infer(image_input, audience, keyword, protagonist):
    global story_intro, current_story, current_image_url, chapter_num, chapters
    chapter_num = 1
    chapters = []
    gr.Info('Calling CLIP Interrogator, 이미지를 해석하고 있습니다...')
    clipi_result = clipi_client.predict(image_input, "best", 4, api_name="/clipi2")[0]
    story_intro = f"""
    # Illustrated Tales
    ## Created by [Sigkawat Pengnoo](https://flowgpt.com/prompt/qzv2D3OvHkzkfSE4rQCqv) at FlowGPT
    Keyword: {keyword}
    Protagonist: {protagonist}
    한국어로 답변해줘.
    STORY : "{{ {clipi_result} }}"
    Let's begin with Chapter 1!
    """
    current_story = clipi_result
    return next_chapter(audience, keyword, protagonist)

css = """
#col-container {max-width: 910px; margin-left: auto; margin-right: auto;}
a {text-decoration-line: underline; font-weight: 600;}
"""
with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
            <h1 style="text-align: center">Illustrated Tales - Korean</h1>
            <p style="text-align: center">이미지를 업로드하세요, ChatGPT를 통해 한국어로 이야기와 그림을 만들어 줍니다!</p>
            """
        )
        with gr.Row():
            with gr.Column():
                image_in = gr.Image(label="이미지 입력", type="filepath", elem_id="image-in", height=420)
                audience = gr.Radio(label="대상", choices=["Children", "Adult"], value="Children")
                keyword_in = gr.Textbox(label="핵심 키워드")
                protagonist_in = gr.Textbox(label="주인공")
                submit_btn = gr.Button('이야기와 그림을 만들어 주세요')
                next_chapter_btn = gr.Button('다음 이야기')
                save_pdf_btn = gr.Button('PDF로 저장하기')
            with gr.Column(elem_id="component-10"):
                chapter_story = gr.Markdown(label="이야기", elem_id="chapter_story")
                chapter_image = gr.Image(label="그림", elem_id="chapter_image")

        table_markdown = gr.Markdown(label="이야기와 그림 정리", elem_id="table_markdown", width="100%")

    submit_btn.click(fn=infer, inputs=[image_in, audience, keyword_in, protagonist_in], outputs=[chapter_story, chapter_image, table_markdown])
    next_chapter_btn.click(fn=lambda: next_chapter(audience=audience.value, keyword=keyword_in.value, protagonist=protagonist_in.value), outputs=[chapter_story, chapter_image, table_markdown])
    save_pdf_btn.click(fn=save_as_pdf, outputs=gr.File(label="PDF 다운로드"))

    demo.queue(max_size=12).launch()