import os
# 環境変数からDIFY_BASE_URLとDIFY_API_KEY_MYWORKFLOWを取得
DIFY_BASE_URL = os.environ.get("DIFY_BASE_URL", "")
DIFY_API_KEY_MYWORKFLOW = os.environ.get("DIFY_API_KEY_MYWORKFLOW", "")

import gradio as gr
import requests
from langchain_community.document_loaders import UnstructuredPDFLoader
import json

def run_workflow(message):
    try:
        file = message['files'][0]
        text_message = message['text']
        
        # PDFファイルが選択されているかチェック
        if not file:
            return "PDFファイルを選択してください。", ""
        
        # PDFファイルをロードしてテキストを抽出
        loader = UnstructuredPDFLoader(file)
        data = loader.load()
        raw_text = data[0].page_content
        
        # APIリクエストのための入力データを準備
        inputs = {
            "url": "",
            "knowledge": raw_text
        }
        
        yield raw_text, "loading...", {}
        
        # APIエンドポイントURL
        url = DIFY_BASE_URL + "/workflows/run"
        
        # APIリクエストのヘッダー
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {DIFY_API_KEY_MYWORKFLOW}"
        }
        
        # APIリクエストのデータ
        data = {
            "inputs": inputs,
            "query": "",
            "response_mode": "streaming",
            "user": "abc_123",
        }
        
        # APIにリクエストを送信
        response = requests.post(url, headers=headers, json=data, stream=True)
        response.raise_for_status()
        
        assistant_message = ""
        outputs = {}
        
        # APIレスポンスのチャンク処理
        for chunk in response.iter_lines(delimiter=b"\n\n"):
            if chunk:
                chunk_data = chunk.decode("utf-8").strip()
                if chunk_data.startswith("data:"):
                    json_data = chunk_data[6:]  # "data: "を取り除く
                    if json_data:
                        result = json.loads(json_data)
                        if result.get("event") == "text_chunk":
                            answer = result.get("data", "").get("text", "")
                            assistant_message += str(answer)
                            yield raw_text, assistant_message, result.get("data", "")
                        elif result.get("event") == "workflow_finished":
                            outputs = result.get('data', "")
                            yield raw_text, assistant_message, outputs
        
    except Exception as e:
        error_message = str(e)
        print(f"Error: {error_message}")
        return "error", error_message, {}

# Gradioインターフェイスの設定
iface = gr.Interface(
    fn=run_workflow,
    inputs=[gr.MultimodalTextbox(label="PDFファイルをアップロード", file_types=[".pdf"], interactive=True)],
    outputs=[
        gr.Textbox(label="生テキスト", show_copy_button=True, max_lines=5),
        gr.Markdown(),
        gr.JSON()
    ],
    title="PDF to Dify Workflow",
    description="PDFファイルを入力すると、Dify APIのワークフローによって処理された結果が表示されます。",
    article="""
    
    © 2024 @tregu0458. All rights reserved.
    
    [利用規約はこちら](https://huggingface.co/spaces/tregu0458/pdf_2_dify_workflow/blob/main/README.md)
    
    ## 使用コンポーネント
    - dify
    - gradio
    - langchain_community.document_loaders
    
    ## 今回のworkflowの仕様
    ### 入力
    - url
    - knowledge
    ### 出力
    - result
    - row_content
    - url
    ### LLM
    - gemini-1.5-flash
    ```
        PDFファイルを入力として受け取り、Dify APIのワークフローを使用してファイルを処理し、結果を返す関数。
        Args:
            message (dict): 入力メッセージ。以下のキーを含む辞書。
                - 'files' (list): アップロードされたPDFファイルのリスト。
                - 'text' (str): テキストメッセージ。
        Yields:
            tuple: 以下の要素を含むタプル。
                - raw_text (str): PDFファイルから抽出された生テキスト。
                - assistant_message (str): アシスタントからのメッセージ。
                - outputs (dict): APIレスポンスのデータ。
        Returns:
            tuple: 以下の要素を含むタプル。
                - status (str): 処理の状態。"error" または "" (空文字列)。
                - error_message (str): エラーメッセージ (エラーが発生した場合)。
                - data (dict): APIレスポンスのデータ。
        Raises:
            Exception: 処理中にエラーが発生した場合。
        Notes:
            - 関数は非同期的に実行され、処理の進行状況に応じて段階的に結果を返す。
            - `yield` を使用して、処理の途中経過を表示しながら、最終的な結果を返す。
    ```
    """
)

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