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)

GEN_RANGE = {
    "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]
}

generation = gr.Dropdown(
            [f"{k}세대" for k in range(1, 10)], value="1세대", label="포켓몬 세대", info="원하는 포켓몬 세대를 선택하세요."
        )

download = gr.File(label="파란 글씨를 눌러서 다운로드 받으세요.")
text = gr.DataFrame(label="포켓몬 리스트")

def write_pdf(gen, progress=gr.Progress()):
    filename = f'포켓몬{gen}.pdf'
    pdfmetrics.registerFont(TTFont("나눔고딕", "NanumGothic.ttf"))
    # 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),
    ])

    story = []
    data_dict = []
    start, end = GEN_RANGE[gen]
    for k in progress.tqdm(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']

        data_dict.append(
            dict(이름=name, No=number, 타입='+'.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)
        table.setStyle(style)
        story.append(table)

    doc.build(story)
    df = pd.DataFrame(data_dict)
    return df[['No', '이름', '타입']], filename

demo = gr.Interface(write_pdf, generation, [text, download], title="대치동 포켓몬 도감 생성기",
                    description="원하는 포켓몬 세대를 선택하고, 다운로드를 눌러주세요.")
demo.queue(concurrency_count=3)
demo.launch()