import gradio as gr
from deepface import DeepFace
import cv2
import numpy as np
import os
import shutil

# --- 設定 ---
DB_PATH = "my_db"

# 起動時にDBフォルダをリセット（Hugging Faceの環境をクリーンに保つため）
if os.path.exists(DB_PATH):
    shutil.rmtree(DB_PATH)
os.makedirs(DB_PATH, exist_ok=True)

# --- 関数定義 ---

def register_face(image, name):
    """画像をDBフォルダに保存する"""
    if image is None or name.strip() == "":
        return "画像と名前を入力してください。"
    
    try:
        # OpenCV形式に変換して保存
        image_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
        # ここを修正：os.path.join です（スラッシュは不要）
        file_path = os.path.join(DB_PATH, f"{name}.jpg")
        cv2.imwrite(file_path, image_bgr)
        return f"「{name}」さんを登録しました！"
    except Exception as e:
        return f"エラーが発生しました: {str(e)}"

def recognize_face(target_image):
    """画像から登録済みの顔を探す"""
    if target_image is None:
        return None, "画像をアップロードしてください。"
    
    if not os.listdir(DB_PATH):
        return target_image, "まだ誰も登録されていません。「推し登録」タブで登録してください。"

    try:
        # 認識の実行
        results = DeepFace.find(img_path=target_image, db_path=DB_PATH, enforce_detection=False, silent=True)
        
        output_image = target_image.copy()
        found_names = []

        for df in results:
            if df.empty: continue
            for index, row in df.iterrows():
                # 登録名を取得
                full_path = row['identity']
                name = os.path.basename(full_path).split('.')[0]
                found_names.append(name)
                
                # 座標の取得と描画
                x, y, w, h = int(row['source_x']), int(row['source_y']), int(row['source_w']), int(row['source_h'])
                cv2.rectangle(output_image, (x, y), (x+w, y+h), (0, 255, 0), 3)
                cv2.putText(output_image, name, (x, y-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
        
        if not found_names:
            message = "登録された推しは見つかりませんでした。"
        else:
            message = f"発見しました！: {', '.join(set(found_names))}"

        return output_image, message

    except Exception as e:
        return target_image, f"認識エラー: {str(e)}"

# --- UI (Gradio) の構築 ---

with gr.Blocks() as demo:
    gr.Markdown("# 🌟 推し認識AIプロトタイプ")
    gr.Markdown("まずは『推し登録』で写真をアップし、その後『推しを探す』で判定してください。")
    
    with gr.Tabs():
        with gr.TabItem("① 推し登録"):
            with gr.Row():
                reg_input_img = gr.Image(label="推しの写真")
                reg_input_name = gr.Textbox(label="推しの名前")
            reg_btn = gr.Button("データベースに登録")
            reg_status = gr.Textbox(label="登録状況")
            
            reg_btn.click(register_face, inputs=[reg_input_img, reg_input_name], outputs=reg_status)
            
        with gr.TabItem("② 推しを探す"):
            with gr.Row():
                rec_input_img = gr.Image(label="判定したい画像")
                rec_output_img = gr.Image(label="判定結果")
            rec_btn = gr.Button("推しを検索！")
            rec_status = gr.Textbox(label="結果メッセージ")
            
            rec_btn.click(recognize_face, inputs=[rec_input_img], outputs=[rec_output_img, rec_status])

# アプリ起動
if __name__ == "__main__":
    demo.launch()