import time
from dataclasses import dataclass, field
from typing import Optional, Dict, Tuple, List

import gradio as gr
import numpy as np
from PIL import Image

from deepface import DeepFace


# =========================
# 參數（可微調）
# =========================
ANALYZE_EVERY_SEC = 0.8
DOWNSAMPLE_W = 720
EMA_ALPHA = 0.45
CONF_TH = 40.0
SWITCH_CONFIRM_SEC = 2.5
IOU_TH = 0.25

# 建議改為 "opencv" 穩定性較高，若想維持原狀可改回 "mtcnn"
DETECTOR_BACKEND = "opencv" 
ALIGN_FACE = True


# =========================
# 情緒中文 + 分數映射
# =========================
EMO_ZH = {
    "angry": "生氣",
    "disgust": "厭惡",
    "fear": "害怕",
    "happy": "開心",
    "sad": "難過",
    "surprise": "驚訝",
    "neutral": "平靜",
    "unknown": "未知",
}

EMO_TO_SCORE = {
    "angry": 1,
    "disgust": 1,
    "fear": 1,
    "sad": 2,
    "neutral": 3,
    "surprise": 4,
    "happy": 5,
    "unknown": 3,
}

# ### [整合需求 新增] ###
def get_emotion_score(emo: str) -> int:
    return EMO_TO_SCORE.get(emo, 3)
# ######################

SCORE_LABELS = {
    1: "1（心情差）",
    2: "2（不太好）",
    3: "3（普通）",
    4: "4（不錯）",
    5: "5（超棒）",
}
SCORE_CHOICES = [SCORE_LABELS[i] for i in [1, 2, 3, 4, 5]]


# =========================
# 狀態
# =========================
@dataclass
class AppState:
    mode: str = "攝影機辨識"

    running: bool = True
    finished: bool = False
    last_analyze_t: float = 0.0

    # tracking
    track_bbox: Optional[Tuple[int, int, int, int]] = None

    # smoothing
    ema_scores: Optional[np.ndarray] = None
    emo_keys: List[str] = field(default_factory=list)

    # stable/candidate
    stable_emo: str = "unknown"
    candidate_emo: Optional[str] = None
    candidate_since: Optional[float] = None

    # final
    final_emo: Optional[str] = None
    final_conf: Optional[float] = None

    # ### [整合需求 新增] 用來儲存整數分數 ###
    final_score: int = 3  
    # ######################


# =========================
# UI HTML
# =========================
def _hint_html(msg: str) -> str:
    return f"""<div style="border-radius:14px;padding:14px;border:1px dashed #DADADA;background:#FAFAFA;color:#111;">{msg}</div>"""

def _camera_hint_before_start() -> str:
    return _hint_html("請先按「Click to access webcam / 允許」開啟攝影機；接著，請按「錄製（Record）」開始。")

def _camera_hint_done() -> str:
    return _hint_html("已完成辨識，可按「重新攝影機辨識」再次進行辨識。")

def _camera_hint_stopped() -> str:
    return _hint_html("已停止辨識。可按「重新攝影機辨識」再試一次。")

def _upload_hint_before() -> str:
    return _hint_html("請上傳一張清楚的臉部照片（臉越大越準）。")

def _upload_hint_done() -> str:
    return _hint_html("已完成圖片辨識，可按「重新上傳圖片辨識」再次辨識。")

def _upload_hint_stopped() -> str:
    return _hint_html("已停止圖片辨識。可按「重新上傳圖片辨識」再試一次。")

def _result_card_html(emo_key: str, conf: Optional[float]) -> str:
    zh = EMO_ZH.get(emo_key, emo_key)
    score = EMO_TO_SCORE.get(emo_key, 3)
    conf_txt = "" if conf is None else f"{conf:.1f}%"
    score_desc = SCORE_LABELS[score].split("（")[1].rstrip("）")

    return f"""
    <div style="border-radius:16px;padding:20px;border:1px solid #E6E6E6;background:#FFFFFF;color:#111;">
      <div style="font-size:14px;color:#444;margin-bottom:8px;">辨識結果</div>
      <div style="font-size:42px;font-weight:800;line-height:1.1;margin-bottom:14px;color:#000;">{zh}</div>
      <div style="font-size:16px;color:#111;margin-bottom:6px;">心情分數：<b>{score}</b>（{score_desc}）</div>
      <div style="font-size:14px;color:#111;">信心值：{conf_txt}</div>
    </div>
    """

def _score_radio_value(score: int) -> str:
    score = int(score)
    score = min(5, max(1, score))
    return SCORE_LABELS[score]

def _downsample_rgb(frame: np.ndarray, target_w: int) -> np.ndarray:
    if frame is None: return None
    h, w = frame.shape[:2]
    if w <= target_w: return frame
    scale = target_w / w
    new_h = max(1, int(h * scale))
    img = Image.fromarray(frame)
    img = img.resize((target_w, new_h))
    return np.array(img)

def _iou(a, b) -> float:
    ax, ay, aw, ah = a
    bx, by, bw, bh = b
    ax2, ay2 = ax + aw, ay + ah
    bx2, by2 = bx + bw, by + bh
    ix1, iy1 = max(ax, bx), max(ay, by)
    ix2, iy2 = min(ax2, bx2), min(ay2, by2)
    iw, ih = max(0, ix2 - ix1), max(0, iy2 - iy1)
    inter = iw * ih
    if inter <= 0: return 0.0
    union = aw * ah + bw * bh - inter
    return inter / max(1e-6, union)

def _pick_face_by_tracking(faces, prev_bbox):
    if not faces: return None, None
    cand = []
    for f in faces:
        fa = f.get("facial_area", {}) or {}
        x, y, w, h = int(fa.get("x",0)), int(fa.get("y",0)), int(fa.get("w",0)), int(fa.get("h",0))
        if w > 0 and h > 0 and f.get("face") is not None:
            cand.append((f, (x, y, w, h), w * h))
    if not cand: return None, None
    if prev_bbox is not None:
        best, best_i = None, 0.0
        for f, bb, area in cand:
            i = _iou(prev_bbox, bb)
            if i > best_i: best_i = i; best = (f, bb)
        if best is not None and best_i >= IOU_TH: return best[0], best[1]
    cand.sort(key=lambda x: x[2], reverse=True)
    return cand[0][0], cand[0][1]

def _reset_state_for_camera(st: AppState) -> AppState:
    st.running = True; st.finished = False; st.last_analyze_t = 0.0
    st.track_bbox = None; st.ema_scores = None; st.emo_keys = []
    st.stable_emo = "unknown"; st.candidate_emo = None; st.candidate_since = None
    st.final_emo = None; st.final_conf = None
    # ### [整合需求] 重置分數
    st.final_score = 3
    return st

def _reset_state_for_upload(st: AppState) -> AppState:
    st.running = False; st.finished = False; st.last_analyze_t = 0.0
    st.track_bbox = None; st.ema_scores = None; st.emo_keys = []
    st.stable_emo = "unknown"; st.candidate_emo = None; st.candidate_since = None
    st.final_emo = None; st.final_conf = None
    # ### [整合需求] 重置分數
    st.final_score = 3
    return st


# =========================
# Camera streaming
# =========================
def on_stream(frame_rgb: np.ndarray, st: AppState):
    if st.mode != "攝影機辨識":
        return gr.update(), gr.update(), gr.update(), gr.update(value=_upload_hint_before()), st, gr.update(value="重新上傳圖片辨識")

    if st.finished or (not st.running):
        cam_visible = False
        if st.final_emo:
            result_html = _result_card_html(st.final_emo, st.final_conf)
            mood_val = _score_radio_value(EMO_TO_SCORE.get(st.final_emo, 3))
            hint_html = _camera_hint_done()
        else:
            result_html = _hint_html("尚未完成辨識。")
            mood_val = _score_radio_value(3)
            hint_html = _camera_hint_stopped()
        return gr.update(visible=cam_visible), gr.update(value=result_html), gr.update(value=mood_val), gr.update(value=hint_html), st, gr.update(value="重新攝影機辨識")

    if frame_rgb is None:
        return gr.update(visible=True), gr.update(value=_hint_html("等待影像…（尚未開始串流）")), gr.update(value=_score_radio_value(3)), gr.update(value=_camera_hint_before_start()), st, gr.update(value="重新攝影機辨識")

    now = time.time()
    if now - st.last_analyze_t < ANALYZE_EVERY_SEC:
        show = st.stable_emo or "unknown"
        return gr.update(visible=True), gr.update(value=_result_card_html(show, st.final_conf)), gr.update(value=_score_radio_value(EMO_TO_SCORE.get(show, 3))), gr.update(value=_camera_hint_before_start()), st, gr.update(value="重新攝影機辨識")

    st.last_analyze_t = now
    small = _downsample_rgb(frame_rgb, DOWNSAMPLE_W)
    faces = DeepFace.extract_faces(img_path=small, detector_backend=DETECTOR_BACKEND, enforce_detection=False, align=ALIGN_FACE)
    pick, bbox = _pick_face_by_tracking(faces, st.track_bbox)

    if pick is None or bbox is None or pick.get("face") is None:
        show = st.stable_emo or "unknown"
        return gr.update(visible=True), gr.update(value=_result_card_html(show, st.final_conf)), gr.update(value=_score_radio_value(EMO_TO_SCORE.get(show, 3))), gr.update(value=_camera_hint_before_start()), st, gr.update(value="重新攝影機辨識")

    st.track_bbox = bbox
    face = pick["face"]
    if face.dtype != np.uint8: face = np.clip(face * 255.0, 0, 255).astype(np.uint8)

    result = DeepFace.analyze(img_path=face, actions=["emotion"], enforce_detection=False)
    if isinstance(result, list): result = result[0]
    emo_dict = result.get("emotion", None)

    if not isinstance(emo_dict, dict) or len(emo_dict) == 0:
        show = st.stable_emo or "unknown"
        return gr.update(visible=True), gr.update(value=_result_card_html(show, st.final_conf)), gr.update(value=_score_radio_value(EMO_TO_SCORE.get(show, 3))), gr.update(value=_camera_hint_before_start()), st, gr.update(value="重新攝影機辨識")

    top_key = max(emo_dict, key=lambda k: float(emo_dict[k]))
    top_conf = float(emo_dict[top_key])

    if top_conf < CONF_TH:
        show = st.stable_emo or "unknown"
        return gr.update(visible=True), gr.update(value=_result_card_html(show, st.final_conf)), gr.update(value=_score_radio_value(EMO_TO_SCORE.get(show, 3))), gr.update(value=_camera_hint_before_start()), st, gr.update(value="重新攝影機辨識")

    keys = list(emo_dict.keys())
    vec = np.array([float(emo_dict[k]) for k in keys], dtype=np.float32)

    if st.ema_scores is None or st.emo_keys != keys:
        st.ema_scores = vec; st.emo_keys = keys
    else:
        st.ema_scores = EMA_ALPHA * vec + (1.0 - EMA_ALPHA) * st.ema_scores

    stable_idx = int(np.argmax(st.ema_scores))
    stable_key = st.emo_keys[stable_idx]
    st.stable_emo = stable_key

    if st.candidate_emo != stable_key:
        st.candidate_emo = stable_key; st.candidate_since = now
    else:
        if st.candidate_since is None: st.candidate_since = now

    confirmed = (st.candidate_since is not None and (now - st.candidate_since) >= SWITCH_CONFIRM_SEC)

    if confirmed and st.final_emo is None:
        st.final_emo = stable_key; st.final_conf = top_conf
        
        # ### [整合需求] 存入分數 ###
        st.final_score = EMO_TO_SCORE.get(stable_key, 3)
        # #######################
        
        st.running = False; st.finished = True
        return gr.update(visible=False), gr.update(value=_result_card_html(stable_key, top_conf)), gr.update(value=_score_radio_value(EMO_TO_SCORE.get(stable_key, 3))), gr.update(value=_camera_hint_done()), st, gr.update(value="重新攝影機辨識")

    return gr.update(visible=True), gr.update(value=_result_card_html(stable_key, top_conf)), gr.update(value=_score_radio_value(EMO_TO_SCORE.get(stable_key, 3))), gr.update(value=_camera_hint_before_start()), st, gr.update(value="重新攝影機辨識")


# =========================
# Upload image
# =========================
def on_upload(img: np.ndarray, st: AppState):
    if st.mode != "上傳圖片辨識":
        return gr.update(), gr.update(), st, gr.update(value=_camera_hint_before_start()), gr.update(value="重新攝影機辨識")

    st = _reset_state_for_upload(st)
    if img is None:
        return gr.update(value=_upload_hint_before()), gr.update(value=_score_radio_value(3)), st, gr.update(value=_upload_hint_before()), gr.update(value="重新上傳圖片辨識")

    # ### [修正] 圖片格式處理 (解決 PNG/灰階 問題) ###
    if img.ndim == 2:
        img = np.stack([img, img, img], axis=-1)
    if img.shape[-1] == 4:
        img = img[:, :, :3]
    # ###########################################

    small = _downsample_rgb(img, DOWNSAMPLE_W)
    faces = DeepFace.extract_faces(img_path=small, detector_backend=DETECTOR_BACKEND, enforce_detection=False, align=ALIGN_FACE)
    pick, bbox = _pick_face_by_tracking(faces, None)
    
    if pick is None or pick.get("face") is None:
        return gr.update(value=_hint_html("沒有偵測到人臉，請換一張更清楚、臉更大的照片。")), gr.update(value=_score_radio_value(3)), st, gr.update(value=_upload_hint_before()), gr.update(value="重新上傳圖片辨識")

    face = pick["face"]
    if face.dtype != np.uint8: face = np.clip(face * 255.0, 0, 255).astype(np.uint8)

    result = DeepFace.analyze(img_path=face, actions=["emotion"], enforce_detection=False)
    if isinstance(result, list): result = result[0]
    emo_dict = result.get("emotion", None)
    if not isinstance(emo_dict, dict) or len(emo_dict) == 0:
        return gr.update(value=_hint_html("辨識失敗，請換一張照片再試一次。")), gr.update(value=_score_radio_value(3)), st, gr.update(value=_upload_hint_before()), gr.update(value="重新上傳圖片辨識")

    top_key = max(emo_dict, key=lambda k: float(emo_dict[k]))
    top_conf = float(emo_dict[top_key])

    st.final_emo = top_key; st.final_conf = top_conf
    
    # ### [整合需求] 存入分數 ###
    st.final_score = EMO_TO_SCORE.get(top_key, 3)
    # #######################
    
    st.finished = True
    return gr.update(value=_result_card_html(top_key, top_conf)), gr.update(value=_score_radio_value(EMO_TO_SCORE.get(top_key, 3))), st, gr.update(value=_upload_hint_done()), gr.update(value="重新上傳圖片辨識")


# =========================
# Mode / Buttons
# =========================
def on_mode_change(mode: str, st: AppState):
    st.mode = mode
    if mode == "攝影機辨識":
        st = _reset_state_for_camera(st)
        return gr.update(visible=True), gr.update(visible=False, value=None), gr.update(value=_hint_html("等待影像…（尚未開始串流）")), gr.update(value=_score_radio_value(3)), gr.update(value=_camera_hint_before_start()), st, gr.update(value="重新攝影機辨識")
    else:
        st = _reset_state_for_upload(st)
        return gr.update(visible=False), gr.update(visible=True, value=None), gr.update(value=_upload_hint_before()), gr.update(value=_score_radio_value(3)), gr.update(value=_upload_hint_before()), st, gr.update(value="重新上傳圖片辨識")

def on_restart(st: AppState):
    if st.mode == "攝影機辨識":
        st = _reset_state_for_camera(st)
        return gr.update(visible=True), gr.update(value=_hint_html("等待影像…（尚未開始串流）")), gr.update(value=_score_radio_value(3)), gr.update(value=_camera_hint_before_start()), st, gr.update(value="重新攝影機辨識"), gr.update()
    else:
        st = _reset_state_for_upload(st)
        return gr.update(), gr.update(value=_upload_hint_before()), gr.update(value=_score_radio_value(3)), gr.update(value=_upload_hint_before()), st, gr.update(value="重新上傳圖片辨識"), gr.update(value=None)

def on_stop(st: AppState):
    if st.mode == "攝影機辨識":
        st.running = False; st.finished = True
        if st.final_emo is None:
            return gr.update(visible=False), gr.update(value=_hint_html("已停止辨識。")), gr.update(value=_score_radio_value(3)), gr.update(value=_camera_hint_stopped()), st, gr.update(value="重新攝影機辨識"), gr.update()
        return gr.update(visible=False), gr.update(value=_result_card_html(st.final_emo, st.final_conf)), gr.update(value=_score_radio_value(EMO_TO_SCORE.get(st.final_emo, 3))), gr.update(value=_camera_hint_done()), st, gr.update(value="重新攝影機辨識"), gr.update()
    
    st = _reset_state_for_upload(st)
    return gr.update(), gr.update(value=_hint_html("已停止圖片辨識。")), gr.update(value=_score_radio_value(3)), gr.update(value=_upload_hint_stopped()), st, gr.update(value="重新上傳圖片辨識"), gr.update(value=None)

css = """#app_container { max-width: 960px; margin: 0 auto; } .gradio-container { background: #f1f5f9 !important; }"""

if __name__ == "__main__":
    with gr.Blocks(title="Emotion Detector", css=css) as demo:
        # 隊友的單獨測試介面 (略)
        pass 
    # demo.launch()