import gradio as gr
import cohere
import psycopg2
from psycopg2.extras import RealDictCursor
import os
from dotenv import load_dotenv

load_dotenv()

# 環境変数の設定
os.environ["PYTHONIOENCODING"] = "utf-8"

# --- 設定 ---
COHERE_API_KEY = os.getenv("COHERE_API_KEY")
DB_CONFIG = {
    "host": "www.ryhintl.com",
    "dbname": "smair",
    "user": "smairuser",
    "password": "smairuser",
    "port": 10629,
    "connect_timeout": 15
}

co = cohere.Client(COHERE_API_KEY)

def generate_ai_advice(query, results):
    context = ""
    for res in results:
        context += f"- 事例: {res['title']}\n  現象: {res['phenomenon']}\n  原因: {res['cause']}\n  処置: {res['action']}\n  防止策: {res['prevention']}\n\n"

    prompt = f"""
あなたはベテランの工場管理者です。以下の「ユーザーからの相談内容」に対し、提供された「過去の類似事例」を参考にして、現場で今すぐ取るべき行動と、長期的な対策をプロの視点で要約・アドバイスしてください。

# ユーザーからの相談内容:
{query}

# 過去の類似事例:
{context}

# 回答の指針:
1. 過去事例から共通する原因を特定してください。
2. 現場担当者が優先して確認すべき点（チェックポイント）を3つ挙げてください。
3. 再発防止のための改善提案を簡潔に述べてください。
回答は日本語で、現場で読みやすいように箇条書きを活用してください。
"""
    # model名を安定した名称に変更（command-r-plusなど）
    response = co.chat(
        message=prompt,
        model='command-a-03-2025' 
    )
    return response.text

def search_troubles(query):
    if not query.strip():
        return "⚠️ トラブル事象を入力してください。"

    conn = None
    try:
        # 1. ベクトル化
        res = co.embed(
            texts=[query],
            model='embed-multilingual-v3.0',
            input_type='search_query'
        )
        query_embedding = res.embeddings[0]

        # 2. DB接続
        conn = psycopg2.connect(**DB_CONFIG, client_encoding='utf8')
        cur = conn.cursor(cursor_factory=RealDictCursor)
        
        search_query = """
        SELECT id, occurrence_date, process_name, equipment_name, title, 
               phenomenon, cause, action, prevention, severity,
               1 - (embedding <=> %s::vector) AS similarity
        FROM factory_troubles
        ORDER BY similarity DESC
        LIMIT 3;
        """
        cur.execute(search_query, (query_embedding,))
        results = cur.fetchall()
        
        if not results:
            return "該当する事例が見つかりませんでした。"

        # 3. AIによるアドバイス生成
        ai_advice = generate_ai_advice(query, results)

        # 4. フォーマット
        output = "## 🤖 AIからのトラブルシューティング・アドバイス\n\n"
        output += f"{ai_advice}\n\n"
        output += "---\n"
        output += "## 🔍 参考にした過去の類似事例\n\n"
        
        for i, res in enumerate(results):
            output += f"### {i+1}. {res['title']} (類似度: {res['similarity']:.4f})\n"
            output += f"**【処置】**: {res['action']}\n\n"
            output += f"""
<details>
<summary><b>📄 元のデータ（出典詳細）を表示</b></summary>
<div style="background-color: #f9f9f9; padding: 10px; border-radius: 5px; border: 1px solid #ddd; color: #333;">
    <p><b>発生日:</b> {res['occurrence_date']} | <b>工程:</b> {res['process_name']}</p>
    <p><b>現象:</b> {res['phenomenon']}</p>
    <p><b>原因:</b> {res['cause']}</p>
    <p><b>防止策:</b> {res['prevention']}</p>
</div>
</details>\n\n---\n"""
        return output

    except Exception as e:
        return f"### ⚠️ エラーが発生しました\n詳細: `{repr(e)}`"
    finally:
        if conn: conn.close()

# --- Gradio Blocks UI ---
with gr.Blocks(title="AI工場トラブル・コンシェルジュ", css=".gradio-container {background-color: #f5f7f9}") as demo:
    gr.Markdown("# 🛠 AI工場トラブル・コンシェルジュ")
    gr.Markdown("過去の知見に基づき、AIが現場の課題解決に向けたアドバイスを生成します。")
    
    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(
                label="現在のトラブル状況を入力", 
                info="例: コンベアのモーターが高温になっており、焦げたような臭いがする。 加工面の焼け発生している。 末端のシール部からの液漏れがある。",
                lines=3
            )
            # ボタンクリック時に非活性化するように設定可能
            search_btn = gr.Button("🔍 知見を検索してAIアドバイスを受ける", variant="primary")
            
    with gr.Row():
        # ステータス表示用のコンポーネント（任意ですが、標準機能で十分カバーされます）
        output_markdown = gr.Markdown(value="ここに結果が表示されます。")

    # .click の際に、検索中はボタンを無効化し、ローディングを表示する設定
    search_btn.click(
        fn=search_troubles,
        inputs=input_text,
        outputs=output_markdown,
        show_progress="full",    # "full"にすることで画面全体にローディングインジケータを表示
        scroll_to_output=True    # 完了時に結果へスクロール
    )

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