# app.py
import streamlit as st
import cv2
from fish_tracker import FishTracker


def main():
    st.set_page_config(layout="wide")
    st.title("🐟 观赏鱼健康监控系统（笔记本版）")

    # 初始化模型（只加载一次）
    if 'tracker' not in st.session_state:
        st.session_state.tracker = FishTracker()

    tracker = st.session_state.tracker

    # 视频显示区域
    frame_placeholder = st.empty()
    alert_placeholder = st.empty()

    # 打开摄像头
    cap = cv2.VideoCapture(0)
    if not cap.isOpened():
        st.error("❌ 无法打开摄像头，请检查设备权限。")
        return

    cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)

    try:
        while True:
            ret, frame = cap.read()
            if not ret:
                st.warning("摄像头读取失败")
                break

            # 处理帧
            output_frame, alerts = tracker.process_frame(frame)

            # 转为 RGB 显示（Streamlit 要求）
            rgb_frame = cv2.cvtColor(output_frame, cv2.COLOR_BGR2RGB)
            frame_placeholder.image(rgb_frame, channels="RGB", width="stretch")

            # 显示告警
            if alerts:
                alert_text = "\n\n".join(alerts)
                alert_placeholder.warning(alert_text)
            else:
                alert_placeholder.empty()

    except KeyboardInterrupt:
        pass
    finally:
        cap.release()


if __name__ == "__main__":
    main()