import os
os.system("pip install gradio==3.46.0")

import gradio as gr
import json
import pandas as pd
from reportlab.lib import colors
from reportlab.pdfbase import pdfmetrics
from reportlab.pdfbase.ttfonts import TTFont
from reportlab.platypus import Table, TableStyle, Image, SimpleDocTemplate, PageBreak
from reportlab.lib import colors
from reportlab.lib.pagesizes import A4
from reportlab.lib.units import cm

with open('pokemon.json', 'r') as f:
    pokemons = json.load(f)
pokemons_types = ["모든 타입"] + sorted(set([t for poke in pokemons for t in poke['types']]))
GEN_RANGE = {
    "모든 세대": [1, 1017],
    "1세대": [1, 151],
    "2세대": [152, 251],
    "3세대": [252, 386],
    "4세대": [387, 493],
    "5세대": [494, 649],
    "6세대": [650, 721],
    "7세대": [722, 809],
    "8세대": [810, 905],
    "9세대": [906, 1017]
}

with gr.Blocks(title="대치동 포켓몬 도감 생성기 🐙") as block:
    with gr.Row():
        gr.Markdown("# 대치동 포켓몬 도감 생성기 🐙\n- 포켓몬의 세대별, 타입별로 도감을 만들어주는 곳입니다.\n- 자료의 출처는 http://pokemon.fandom.com 입니다.")
    with gr.Row():
        generation = gr.Dropdown(
                    [f"{k}세대" for k in range(1, 10)] + ["모든 세대"], value="1세대", label="🌈포켓몬 세대", info="원하는 포켓몬 세대를 선택하세요."
                )
        poke_types = gr.Dropdown(
                    pokemons_types, value="모든 타입", label="🩻포켓몬 타입", info="원하는 포켓몬 타입을 선택하세요."
                )
        button = gr.Button(value="🎩 도감 생성")
    with gr.Row():
        with gr.Column():
            download = gr.File(label="🔑파란 글씨를 눌러서 다운로드 받으세요.")
            df_view = gr.DataFrame(label="📝포켓몬 리스트", wrap=True, row_count=10)
        with gr.Column():
            gallery = gr.Gallery(columns=3, label="👩🏻‍🎨포켓몬 갤러리", rows=4, height=600)
            report = gr.Markdown("## 타입별 포켓몬 분석")            

    def write_pdf(gen, poke_type):
        filename = f'포켓몬{gen}_{poke_type}.pdf'
        # PDF 문서를 생성합니다.
        pdfmetrics.registerFont(TTFont("나눔고딕", "NanumGothic.ttf"))
        doc = SimpleDocTemplate(filename, pagesize=A4)
        
        # 테이블 스타일을 정의합니다.
        style = TableStyle([
            ('TEXTCOLOR', (0, 1), (-1, -1), colors.gray),
            ('ALIGN', (0, 0), (-1, -1), 'CENTER'),
            ('FONTNAME', (0, 0), (-1, -1), '나눔고딕'),
            ('SIZE', (0, 0), (-1, -1), 50),
            ('BACKGROUND', (0, 1), (-1, 1), colors.white),
            ('GRID', (0, 0), (-1, -1), 1, colors.black),
            ('VALIGN', (0, 0), (-1, -1), 'MIDDLE'),
            ('LEADING', (0, 0), (-1, -1), 14),  # 라인 간격 조정
            ('BOTTOMPADDING', (0, 0), (-1, -1), 60),
        ])

        images = []
        story = []
        data_dict = []
        start, end = GEN_RANGE[gen]
        for k in range(start, end+1):
            name = pokemons[k-1]['name']
            number = pokemons[k-1]['number']
            types = pokemons[k-1]['types']
            image_path = pokemons[k-1]['image_path']

            if poke_type == "모든 타입":
                condition = True
            else:
                condition = poke_type in types
            if condition:
                data_dict.append(
                    dict(이름=name, 번호=number, 타입='+'.join(types))
                )
                url = f"https://huggingface.co/spaces/yoon-gu/pokemon/resolve/main/{image_path}"
                images.append((url, f"{number} {name}({'+'.join(types)})"))

                # PDF에 이미지를 추가합니다. 이미지 파일의 경로를 지정하세요.
                image = Image(image_path, width=13.5*cm, height=13.5*cm)

                # 테이블 데이터를 준비합니다.
                data = [
                    [number],
                    [image],  # 2번째 행에 이미지를 추가할 것입니다.
                    ['+'.join(types)],
                    [name],
                ]
                # 테이블 객체를 생성합니다.
                table = Table(data, colWidths=17.5*cm)
                table.setStyle(style)
                story.append(table)

        doc.build(story)
        df = pd.DataFrame(data_dict)
        analysis = df[['이름', '타입']].groupby(by='타입').count().sort_values(by='이름', ascending=False).to_markdown().replace('이름', '개수')
        return df[['번호', '이름', '타입']], filename, images, analysis

    button.click(write_pdf, [generation, poke_types], [df_view, download, gallery, report])
block.queue(concurrency_count=3)
block.launch()