import gradio as gr
import google.generativeai as genai
import os
import json
import re
from pathlib import Path

def load_prompt(prompt_file):
    """從檔案讀取 prompt.txt 的內容。"""
    try:
      with open(prompt_file, 'r', encoding='utf-8') as f:
        prompt = f.read()
      return prompt
    except FileNotFoundError:
       print(f"FileNotFoundError: Prompt file not found at {prompt_file}")
       return f"Error: Prompt file not found at {prompt_file}"
    except Exception as e:
        print(f"Error loading prompt file: {e}")
        return f"Error: An unexpected error occurred when processing the prompt. Detail: {e}"

def load_dictionary(dictionary_file):
    """從檔案讀取 dictionary.txt 的內容。"""
    try:
        with open(dictionary_file, 'r', encoding='utf-8') as f:
            return json.load(f)
    except FileNotFoundError:
        print(f"FileNotFoundError: Dictionary file not found at {dictionary_file}")
        return f"Error: Dictionary file not found at {dictionary_file}"
    except json.JSONDecodeError as e:
        print(f"JSONDecodeError: {e}")
        return f"Error: Invalid JSON format in dictionary file. JSONDecodeError: {e}"
    except Exception as e:
        print(f"Error loading dictionary JSON: {e}")
        return f"Error: An unexpected error occurred when processing the JSON. Detail: {e}"


def process_article(api_key, article_content, dictionary_file="dictionary.txt", prompt_file="prompt.txt"):
    """處理文章，將結果送給 Google AI 並返回回應。"""
    try:
        genai.configure(api_key=api_key)
        model = genai.GenerativeModel("gemini-1.5-flash")

        # 讀取文章檔案內容
        if not article_content:
            return "Error: No article provided.", "", ""
        
        prompt_template = load_prompt(prompt_file)

        if isinstance(prompt_template, str) and prompt_template.startswith("Error:"):
            return prompt_template, "", ""
        
        dictionary_data = load_dictionary(dictionary_file)

        if isinstance(dictionary_data, str) and dictionary_data.startswith("Error:"):
            return dictionary_data, "", ""

        prompt = prompt_template.format(article=article_content)
        response = model.generate_content(prompt)
        
        token_text = response.text  # 保存原始斷詞結果
        modified_article = response.text
        report = []

        for item in dictionary_data:
            if 'usecase' in item and 'original' in item and 'modified' in item:
                original = item['original']
                modified = item['modified']
                for usecase_word in item['usecase']:
                    while usecase_word in modified_article:
                       modified_article = re.sub(re.escape(usecase_word), usecase_word.replace(original, modified), modified_article, 1) # 只替換第一次出現
                       report.append({"original": original, "modified": modified, "word": usecase_word})


        modified_article = modified_article.replace('^', '')

        #return modified_article, json.dumps(report, ensure_ascii=False, indent=2), token_text  # 返回結果和報告
        return modified_article, json.dumps(report, ensure_ascii=False, indent=2), token_text

    except Exception as e:
        print(f"Error generating content: {e}")
        return f"Error: An unexpected error occurred: {e}", "", ""


iface = gr.Interface(
    fn=process_article,
    inputs=[
        gr.Textbox(label="請貼上 Google AI API 金鑰"),
        gr.Textbox(label="請輸入需修改的文本"),
    ],
    outputs=[
        gr.Textbox(label="處理後的文章"),
        gr.Textbox(label="處理過的詞彙 (JSON)"),
        gr.Textbox(label="斷詞結果"),
    ],
    title="多音字替換",
    description="請輸入文章文本，並獲得 Google AI 的斷詞結果、處理後的文章和詞彙處理報告。",
)

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