import gradio as gr
from openai import OpenAI
import requests
from PIL import Image
from io import BytesIO
import hashlib
import json
import os
from datetime import datetime, timedelta
from functools import wraps
import time
from typing import Generator
import asyncio
from pathlib import Path

# 常量定義
LANGUAGES = ["中文", "英文", "日文", "韓文", "閩南語"]
STORY_TYPES = ["童話", "科幻", "懸疑", "奇幻", "歷史", "愛情"]
IMAGE_STYLES = ["寫實", "卡通", "水彩", "素描", "國畫", "書法"]

# 自定義錯誤類
class APIError(Exception):
    """自定義 API 錯誤"""
    pass

# 緩存系統
class StoryCache:
    def __init__(self, cache_dir="story_cache"):
        self.cache_dir = Path(cache_dir)
        self.cache_dir.mkdir(exist_ok=True)
        
    def _get_cache_key(self, **kwargs):
        """生成緩存鍵"""
        cache_str = json.dumps(kwargs, sort_keys=True)
        return hashlib.md5(cache_str.encode()).hexdigest()
    
    def get(self, **kwargs):
        """獲取緩存的故事"""
        cache_key = self._get_cache_key(**kwargs)
        cache_file = self.cache_dir / f"{cache_key}.json"
        
        if cache_file.exists():
            with cache_file.open('r', encoding='utf-8') as f:
                cached_data = json.load(f)
                # 檢查緩存是否過期（24小時）
                if datetime.fromisoformat(cached_data['timestamp']) + timedelta(hours=24) > datetime.now():
                    return cached_data['content']
        return None
        
    def set(self, content, **kwargs):
        """設置故事緩存"""
        cache_key = self._get_cache_key(**kwargs)
        cache_file = self.cache_dir / f"{cache_key}.json"
        
        cache_data = {
            'content': content,
            'timestamp': datetime.now().isoformat(),
            'metadata': kwargs
        }
        
        with cache_file.open('w', encoding='utf-8') as f:
            json.dump(cache_data, f, ensure_ascii=False, indent=2)

# 初始化緩存
story_cache = StoryCache()

# 錯誤處理裝飾器
def handle_api_error(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except APIError as e:
            return str(e)
        except Exception as e:
            return f"發生錯誤: {str(e)}\n請稍後再試或聯繫支持團隊。"
    return wrapper

def validate_api_key(api_key):
    """驗證 API 密鑰"""
    if not api_key or len(api_key.strip()) < 10:
        raise APIError("請提供有效的 API 密鑰")
    return api_key.strip()

@handle_api_error
def generate_story(prompt, story_type, language, api_key):
    """生成故事"""
    api_key = validate_api_key(api_key)
    client = OpenAI(api_key=api_key)
    
    system_prompt = f"""你是一位專業的{story_type}故事作家。請遵循以下準則：
    1. 使用{language}創作
    2. 故事結構需包含：開頭、發展、高潮、結局
    3. 根據故事類型添加相應的元素和氛圍
    4. 確保故事適合所有年齡層
    5. 字數控制在1000字左右
    """
    
    try:
        response = client.chat.completions.create(
            model="gpt-3.5-turbo",
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": f"寫一個關於{prompt}的{story_type}短篇故事"}
            ],
            temperature=0.7,
            max_tokens=2000
        )
        return response.choices[0].message.content
    except Exception as e:
        raise APIError(f"故事生成失敗: {str(e)}")

@handle_api_error
def translate_story(story, target_language, api_key):
    """翻譯故事"""
    client = OpenAI(api_key=api_key)
    try:
        response = client.chat.completions.create(
            model="gpt-3.5-turbo",
            messages=[
                {"role": "system", "content": f"你是一位專業翻譯。請將以下故事翻譯成{target_language}，保持原文的風格和感情。"},
                {"role": "user", "content": story}
            ],
            temperature=0.3
        )
        return response.choices[0].message.content
    except Exception as e:
        raise APIError(f"翻譯錯誤: {str(e)}")

@handle_api_error
def generate_image(prompt, style, api_key):
    """生成圖片"""
    client = OpenAI(api_key=api_key)
    try:
        response = client.images.generate(
            prompt=f"以{style}風格為一個關於{prompt}的故事創作插圖。畫面需要精美細緻，構圖完整。",
            size="1024x1024",
            n=1,
            quality="standard"
        )
        image_url = response.data[0].url
        image_response = requests.get(image_url)
        return Image.open(BytesIO(image_response.content))
    except Exception as e:
        raise APIError(f"圖片生成失敗: {str(e)}")

def process_with_progress(prompt, story_type, original_language, target_language, image_style, api_key) -> Generator:
    """使用生成器提供進度更新"""
    steps = [
        ("正在構思故事...", 0.2),
        ("正在生成故事內容...", 0.4),
        ("正在生成插圖...", 0.7),
        ("正在翻譯故事...", 0.9),
        ("完成！", 1.0)
    ]
    
    original_story = None
    translated_story = None
    image = None
    
    try:
        for message, progress in steps:
            yield {"progress": progress, "message": message, "status": "running"}
            
            if progress == 0.4:
                # 檢查緩存
                cache_key = {
                    'prompt': prompt,
                    'story_type': story_type,
                    'language': original_language
                }
                original_story = story_cache.get(**cache_key)
                if not original_story:
                    original_story = generate_story(prompt, story_type, original_language, api_key)
                    story_cache.set(original_story, **cache_key)
                    
            elif progress == 0.7:
                image = generate_image(prompt, image_style, api_key)
                
            elif progress == 0.9 and target_language != original_language:
                translated_story = translate_story(original_story, target_language, api_key)
        
        yield {
            "progress": 1.0,
            "message": "生成完成！",
            "status": "complete",
            "results": {
                "original_story": original_story,
                "translated_story": translated_story if target_language != original_language else "",
                "image": image
            }
        }
    except Exception as e:
        yield {
            "progress": 0,
            "message": f"發生錯誤: {str(e)}",
            "status": "error"
        }

def save_story(story, filename="story.txt"):
    """保存故事"""
    try:
        with open(filename, 'w', encoding='utf-8') as f:
            f.write(story)
        return f"故事已保存為 {filename}"
    except Exception as e:
        return f"故事保存失敗: {str(e)}"

def save_image(image, filename="illustration.png"):
    """保存圖片"""
    try:
        if image is not None:
            image.save(filename)
            return f"插圖已保存為 {filename}"
        return "沒有可保存的插圖"
    except Exception as e:
        return f"插圖保存失敗: {str(e)}"

def export_to_pdf(story, translated_story, image, filename="story_book.pdf"):
    """導出為PDF"""
    try:
        from reportlab.pdfgen import canvas
        from reportlab.lib.pagesizes import A4
        from reportlab.pdfbase import pdfmetrics
        from reportlab.pdfbase.ttfonts import TTFont
        
        # 註冊中文字體（需要自行提供字體文件）
        try:
            pdfmetrics.registerFont(TTFont('ChineseFont', 'path/to/chinese/font.ttf'))
        except:
            pass  # 如果沒有中文字體，使用默認字體
        
        c = canvas.Canvas(filename, pagesize=A4)
        width, height = A4
        
        # 添加標題
        c.setFont('Helvetica-Bold', 24)
        c.drawCentredString(width/2, height-50, "創意故事")
        
        # 添加原文
        c.setFont('ChineseFont' if 'ChineseFont' in pdfmetrics.getRegisteredFontNames() else 'Helvetica', 12)
        text_object = c.beginText()
        text_object.setTextOrigin(50, height-100)
        for line in story.split('\n'):
            text_object.textLine(line)
        c.drawText(text_object)
        
        # 添加翻譯
        if translated_story:
            c.showPage()  # 新頁面
            text_object = c.beginText()
            text_object.setTextOrigin(50, height-50)
            for line in translated_story.split('\n'):
                text_object.textLine(line)
            c.drawText(text_object)
        
        # 添加圖片
        if image:
            c.showPage()
            image_path = "temp_illustration.png"
            image.save(image_path)
            c.drawImage(image_path, 50, height-500, width=400, height=400)
            os.remove(image_path)
        
        c.save()
        return f"已導出為 PDF: {filename}"
    except Exception as e:
        return f"PDF 導出失敗: {str(e)}"

# Gradio 介面
with gr.Blocks(css="""
    .container { max-width: 900px; margin: auto; padding: 20px; }
    .output-box { border: 1px solid #ddd; padding: 10px; border-radius: 5px; }
    .btn { background-color: #4CAF50; color: white; padding: 10px 20px; border-radius: 5px; border: none; }
    .btn:hover { background-color: #45a049; }
""") as demo:
    gr.Markdown("# 🌟 創意故事與插畫生成器 🎨")
    
    with gr.Row():
        with gr.Column(scale=2):
            prompt = gr.Textbox(label="✍️ 故事靈感", placeholder="例如：一隻會說話的貓咪")
        with gr.Column(scale=1):
            api_key = gr.Textbox(label="🔑 OpenAI API 密鑰", type="password")

    with gr.Row():
        with gr.Column():
            story_type = gr.Dropdown(STORY_TYPES, label="📚 故事類型", value="童話")
            original_language = gr.Dropdown(LANGUAGES, label="🗣️ 原始語言", value="中文")
        with gr.Column():
            image_style = gr.Dropdown(IMAGE_STYLES, label="🖌️ 插圖風格", value="水彩")
            target_language = gr.Dropdown(LANGUAGES, label="🌐 翻譯語言", value="英文")

    # 添加進度顯示
    status = gr.Markdown("準備就緒")
    progress = gr.Progress()
    
    generate_btn = gr.Button("🚀 開始創作", elem_classes="btn")

    with gr.Row():
        with gr.Column():
            original_story = gr.Textbox(label="📖 原始故事", elem_classes="output-box")
            translated_story = gr.Textbox(label="🔄 翻譯後的故事", elem_classes="output-box")
        with gr.Column():
            image_output = gr.Image(label="🖼️ 生成的插圖")

    with gr.Row():
        save_story_btn = gr.Button("💾 保存故事", elem_classes="btn")
        save_image_btn = gr.Button("💾 保存插圖", elem_classes="btn")
        export_pdf_btn = gr.Button("📑 導出PDF", elem_classes="btn")

    save_status = gr.Textbox(label="保存狀態")

    def process_with_ui_updates(prompt, story_type, original_language, target_language, image_style, api_key):
        for update in process_with_progress(prompt, story_type, original_language, target_language, image_style, api_key):
            status.value = update["message"]
            if update["status"] == "complete":
                return (
                    update["results"]["original_story"],
                    update["results"]["translated_story"],
                    update["results"]["image"]
                )
            elif update["status"] == "error":
                return update["message"], "", None

    # 事件處理
    generate_btn.click(
        fn=process_with_ui_updates,
        inputs=[prompt, story_type, original_language, target_language, image_style, api_key],
        outputs=[original_story, translated_story, image_output]
    )

    save_story_btn.click(
        fn=save_story,
        inputs=[original_story],
        outputs=[save_status]
    )

    save_image_btn.click(
        fn=save_image,
        inputs=[image_output],
        outputs=[save_status]
    )

    export_pdf_btn.click(
        fn=export_to_pdf,
        inputs=[original_story, translated_story, image_output],
        outputs=[save_status]
    )

    gr.Markdown("""
    ### 📌 使用說明
    1. 輸入你的創意故事靈感
    2. 選擇故事類型、語言和插圖風格""")

if __name__ == "__main__":
    demo.launch()