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

# --- 設定 ---
DB_PATH = "my_db"
if not os.path.exists(DB_PATH):
    os.makedirs(DB_PATH, exist_ok=True)

# 処理を間引くためのカウンター
frame_count = 0
last_status = "Waiting..."

# --- 関数定義 ---

def register_face(image, name):
    """画像を1枚だけDBに保存する"""
    if image is None or name.strip() == "":
        return "名前を入力してください。"
    try:
        # DBを一度空にして、常に「最新の1枚」とだけ比較するようにする（超軽量化）
        for f in os.listdir(DB_PATH):
            os.remove(os.path.join(DB_PATH, f))
            
        safe_name = "".join([c for c in name if c.isalnum()])
        file_path = os.path.join(DB_PATH, f"{safe_name}.jpg")
        cv2.imwrite(file_path, cv2.cvtColor(image, cv2.COLOR_RGB2BGR))
        return f"「{safe_name}」をターゲットに設定しました！"
    except Exception as e:
        return f"エラー: {str(e)}"

def track_oshi(frame):
    """Webカメラの映像を解析（1対1の超軽量比較）"""
    if frame is None: return None
    
    global frame_count, last_status
    frame_count += 1
    output_frame = frame.copy()
    
    # 画面に情報を表示
    cv2.putText(output_frame, f"Frame: {frame_count}", (10, 30), 
                cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 0), 2)
    cv2.putText(output_frame, f"Status: {last_status}", (10, 60), 
                cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 255), 2)

    # 30フレーム（数秒〜10秒に一度）だけ解析を行う
    if frame_count % 30 == 0:
        db_images = [f for f in os.listdir(DB_PATH) if f.endswith(".jpg")]
        if db_images:
            last_status = "Analyzing..."
            target_path = os.path.join(DB_PATH, db_images[0])
            try:
                # findではなく、最も軽い「verify」を使用
                result = DeepFace.verify(
                    img1_path = frame,
                    img2_path = target_path,
                    enforce_detection = False,
                    detector_backend = 'opencv',
                    model_name = 'VGG-Face', # 最も標準的なモデル
                    silent = True
                )
                
                if result["verified"]:
                    last_status = "TARGET FOUND!"
                    # 画面中央に大きく表示
                    cv2.putText(output_frame, "MATCH!!!", (100, 200), 
                                cv2.FONT_HERSHEY_SIMPLEX, 2.0, (0, 255, 0), 5)
                else:
                    last_status = "Not Match"
            except Exception as e:
                last_status = "Scan Error"
    
    return output_frame

# --- UI構築 ---

with gr.Blocks() as demo:
    gr.Markdown("# 🎥 超軽量版・推し認識プロトタイプ")
    gr.Markdown("Hugging Faceの無料CPUでも動くよう、解析頻度を極限まで落としています。")
    
    with gr.Tabs():
        with gr.TabItem("① ターゲット登録"):
            reg_in = gr.Image(label="あなたの顔写真")
            reg_name = gr.Textbox(label="名前（英数字）", value="Me")
            reg_btn = gr.Button("この人をターゲットに設定")
            reg_out = gr.Textbox(label="ログ")
            reg_btn.click(register_face, inputs=[reg_in, reg_name], outputs=reg_out)

        with gr.TabItem("② リアルタイム検証"):
            input_video = gr.Image(sources=["webcam"], streaming=True)
            output_video = gr.Image(label="解析結果（数秒に1回更新）")
            input_video.stream(track_oshi, inputs=[input_video], outputs=[output_video])

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