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

load_dotenv(verbose=True)

# --- 設定 ---
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 register_trouble(date, process, equipment, title, phenomenon, cause, action, prevention, severity):
    if not title or not phenomenon:
        return "⚠️ タイトルと現象は必須入力項目です。"

    conn = None
    try:
        # 1. 登録用テキストを結合してベクトル化 (Embedding生成)
        # 検索精度を高めるため、主要な情報をすべて含める
        input_text = f"【{process}】{title} 現象:{phenomenon} 原因:{cause} 処置:{action} 防止策:{prevention}"
        
        res = co.embed(
            texts=[input_text],
            model='embed-multilingual-v3.0',
            input_type='search_document' # 登録時は document を指定
        )
        embedding = res.embeddings[0]

        # 2. DB接続とインサート
        conn = psycopg2.connect(**DB_CONFIG, client_encoding='utf8')
        cur = conn.cursor()
        
        insert_query = """
        INSERT INTO public.factory_troubles (
            occurrence_date, process_name, equipment_name, title, 
            phenomenon, cause, action, prevention, severity, embedding
        ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
        """
        
        cur.execute(insert_query, (
            date, process, equipment, title, 
            phenomenon, cause, action, prevention, int(severity), 
            embedding
        ))
        
        conn.commit()
        return f"✅ 登録完了: 「{title}」をナレッジベースに保存しました。"

    except Exception as e:
        return f"### ⚠️ 登録エラー\n詳細: `{repr(e)}`"
    finally:
        if conn: conn.close()

# --- 検索関数（前回までの機能） ---
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)
        
        # Citationのために全カラムを取得するように変更
        cur.execute("""
            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
        """, (query_embedding,))
        results = cur.fetchall()
        
        if not results:
            return "該当する事例が見つかりませんでした。"

        # 3. 表示用にフォーマット
        output = f"## 🔍 類似事例の検索結果（クエリ: {query}）\n\n"
        
        for i, r in enumerate(results):
            # メインの回答部分
            output += f"### {i+1}. {r['title']} (類似度: {r['similarity']:.3f})\n"
            output += f"**✅ 推奨される処置:**\n{r['action']}\n\n"
            output += f"**⚠️ 今後の防止策:**\n{r['prevention']}\n\n"
            
            # --- Citation (出典詳細情報) ---
            output += f"""
<details>
<summary><b>📄 この事例の出典詳細を表示 (ID: {r['id']})</b></summary>
<div style="background-color: #f9f9f9; padding: 15px; border-radius: 8px; border: 1px solid #ddd; color: #333; margin-top: 10px;">
    <p><b>📅 発生日:</b> {r['occurrence_date']} | <b>🏢 工程:</b> {r['process_name']} | <b>⚙️ 設備:</b> {r['equipment_name']}</p>
    <p><b>🚨 重要度:</b> {r['severity']} / 5</p>
    <hr style="border: 0.5px solid #ccc;">
    <p><b>【元の現象】</b><br>{r['phenomenon']}</p>
    <p><b>【推定原因】</b><br>{r['cause']}</p>
</div>
</details>
\n"""
            output += "\n---\n\n"
            
        return output

    except Exception as e:
        return f"### ⚠️ 検索エラー\n詳細: `{repr(e)}`"
    finally:
        if conn: conn.close()

# --- Gradio UI 構成 ---
with gr.Blocks(title="工場トラブル・ナレッジ・プラットフォーム", css=".gradio-container {background-color: #f0f2f5}") as demo:
    gr.Markdown("# 🛠 工場トラブル・ナレッジ・プラットフォーム")
    
    with gr.Tab("事例を検索する"):
        with gr.Column():
            search_input = gr.Textbox(label="状況を入力して過去の知見を検索", placeholder="例: モーターの異音...")
            search_btn = gr.Button("検索実行", variant="primary")
            search_output = gr.Markdown()
            
            search_btn.click(fn=search_troubles, inputs=search_input, outputs=search_output)

    with gr.Tab("新規事例を登録する"):
        gr.Markdown("### 📝 新しいトラブル報告の入力")
        with gr.Row():
            date_input = gr.Textbox(label="発生日", value=datetime.now().strftime('%Y-%m-%d'))
            process_input = gr.Dropdown(label="工程名", choices=["露光", "エッチング", "成膜", "プレス", "溶接", "切削", "充填", "包装", "組立", "エッチング", "成膜", "調合", "イオン注入", "洗浄", "CMP", "熱処理", "酸洗浄", "スパッター", "封止", "殺菌", "包装", "調合", "搬送", "テスト", "共通"])
            severity_input = gr.Slider(label="重要度 (1-5)", minimum=1, maximum=5, step=1, value=3)
        
        equip_input = gr.Textbox(label="設備名")
        title_input = gr.Textbox(label="トラブル件名 (Title)")
        phenom_input = gr.TextArea(label="現象 (Phenomenon)")
        cause_input = gr.TextArea(label="原因 (Cause)")
        action_input = gr.TextArea(label="処置 (Action)")
        prevent_input = gr.TextArea(label="防止策 (Prevention)")
        
        reg_btn = gr.Button("🚀 データベースに登録", variant="primary")
        reg_output = gr.Markdown()

        reg_btn.click(
            fn=register_trouble,
            inputs=[date_input, process_input, equip_input, title_input, phenom_input, cause_input, action_input, prevent_input, severity_input],
            outputs=reg_output,
            show_progress="full"
        )

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