UI.py

import os
import sys
import cv2
import requests # 추가
import json     # 추가
import base64   # 추가
import threading # 추가 (UI 멈춤 방지용)
from PyQt6.QtWidgets import (QApplication, QWidget, QVBoxLayout, QHBoxLayout, 
                             QLabel, QPushButton, QFrame, QMessageBox)
from PyQt6.QtCore import Qt, QTimer
from PyQt6.QtGui import QImage, QPixmap

from sensor_manager import SensorManager
from motor_manger import MotorManager
from prediction import run_detection_once
from dotenv import load_dotenv

from http_request import send_trash_result, process_and_send_backend
load_dotenv()
# .env에서 서버 주소 로드

SERIAL_PORT = os.getenv('SERIAL_PORT')
sensor = SensorManager(port=SERIAL_PORT)
sensor.connect()

motor = MotorManager(port=SERIAL_PORT)
motor.connect()




# 기존 StatusRow 클래스 (생략 가능, 그대로 유지)
class StatusRow(QFrame):
    def __init__(self, dot_color, text, percent):
        super().__init__()
        self.setFixedHeight(60)
        self.setStyleSheet(f"""
            StatusRow {{
                background-color: rgba(255, 255, 255, 0.2);
                border-radius: 15px;
            }}
        """)
        layout = QHBoxLayout()
        layout.setContentsMargins(15, 0, 15, 0)
        self.setLayout(layout)
        self.dot = QLabel()
        self.dot.setFixedSize(15, 15)
        self.dot.setStyleSheet(f"background-color: {dot_color}; border-radius: 7px; border: none;")
        self.label = QLabel(text)
        self.label.setStyleSheet("background-color: transparent; color: white; font-size: 16px; font-weight: bold; border: none;")
        self.percent = QLabel(percent)
        self.percent.setStyleSheet("background-color: transparent; color: white; font-size: 18px; font-weight: 900; border: none;")
        layout.addWidget(self.dot)
        layout.addSpacing(10)
        layout.addWidget(self.label)
        layout.addStretch()
        layout.addWidget(self.percent)


class WasteApp(QWidget):
    def __init__(self):
        super().__init__()
        self.cap = None
        self.is_processing = False 
        self.timer = QTimer()
        self.timer.timeout.connect(self.update_frame)
        self.initUI()
        QTimer.singleShot(100, self.initialize_camera)

    def initUI(self):
        # ... (기존 UI 코드와 동일) ...
        self.setWindowTitle('Smart Waste Bin (Landscape)')
        self.setFixedSize(800, 400)
        self.setStyleSheet("background-color: #00B140; font-family: sans-serif;")

        main_layout = QHBoxLayout()
        main_layout.setContentsMargins(30, 30, 30, 30)
        main_layout.setSpacing(30)
        self.setLayout(main_layout)

        # [좌측]
        left_layout = QVBoxLayout()
        self.video_container = QFrame()
        self.video_container.setFixedSize(360, 270)
        self.video_container.setStyleSheet("background-color: rgba(0,0,0,0.1); border-radius: 15px;")
        
        video_layout = QVBoxLayout(self.video_container)
        video_layout.setContentsMargins(0,0,0,0)
        video_layout.setAlignment(Qt.AlignmentFlag.AlignCenter)

        self.top_btn = QPushButton("↻")
        self.top_btn.setFixedSize(100, 100)
        self.top_btn.setStyleSheet("""
            QPushButton {
                background-color: white;
                border-radius: 50px;
                color: #00B140;
                font-size: 50px;
                font-weight: bold;
                border: none;
            }
            QPushButton:pressed { background-color: #E0E0E0; }
        """)
        self.top_btn.clicked.connect(self.start_process)

        self.video_label = QLabel()
        self.video_label.setFixedSize(320, 240)
        self.video_label.setStyleSheet("background-color: black; border-radius: 10px;")
        self.video_label.hide()

        video_layout.addWidget(self.top_btn)
        video_layout.addWidget(self.video_label)
        
        self.subtitle_label = QLabel("카메라 초기화 중...")
        self.subtitle_label.setAlignment(Qt.AlignmentFlag.AlignCenter)
        self.subtitle_label.setStyleSheet("color: #D0F0D0; font-size: 16px; margin-top: 10px;")

        left_layout.addStretch()
        left_layout.addWidget(self.video_container)
        left_layout.addWidget(self.subtitle_label)
        left_layout.addStretch()

        # [우측]
        right_layout = QVBoxLayout()
        self.title_label = QLabel("압축 실행")
        self.title_label.setStyleSheet("color: white; font-size: 36px; font-weight: 900;")
        
        right_layout.addWidget(self.title_label)
        right_layout.addSpacing(20)
        
        self.general_row = StatusRow("#9E9E9E", "일반쓰레기", "67%")
        self.can_row = StatusRow("#FFC107", "캔", "82%")
        self.plastic_row = StatusRow("#42A5F5", "플라스틱", "45%")

        right_layout.addWidget(self.general_row)
        right_layout.addWidget(self.can_row)
        right_layout.addWidget(self.plastic_row)
        right_layout.addStretch()

        main_layout.addLayout(left_layout, 1)
        main_layout.addLayout(right_layout, 1)

    def initialize_camera(self):
        # ... (기존 카메라 연결 코드와 동일) ...
        pipeline = (
            "libcamerasrc ! "
            "video/x-raw, width=640, height=480, framerate=30/1 ! "
            "videoconvert ! "
            "video/x-raw, format=BGR ! "
            "appsink drop=1"
        )
        try:
            self.cap = cv2.VideoCapture(pipeline, cv2.CAP_GSTREAMER)
            if not self.cap.isOpened():
                print("GStreamer 실패, V4L2 시도")
                self.cap = cv2.VideoCapture(0, cv2.CAP_V4L2)
            
            if self.cap.isOpened():
                self.subtitle_label.setText("버튼을 눌러 압축 시작")
            else:
                self.subtitle_label.setText("카메라 연결 실패")
        except Exception as e:
            print(f"카메라 에러: {e}")
            self.subtitle_label.setText("카메라 오류")

    def start_process(self):
        """
        [STEP 1] 버튼 클릭 -> [STEP 2] 분류 -> [백엔드 전송] -> [STEP 3] 서브모터
        """
        if self.is_processing:
            return
        
        if self.cap is None or not self.cap.isOpened():
            print("카메라 재연결 시도...")
            self.initialize_camera()
            if self.cap is None or not self.cap.isOpened():
                self.subtitle_label.setText("카메라 연결 실패")
                return

        self.is_processing = True
        self.top_btn.hide()
        self.video_label.show()
        
        self.subtitle_label.setText("이미지 분류 중...")
        QApplication.processEvents()
        
        detected_label = ""
        current_frame = None

        if self.cap and self.cap.isOpened():
            ret, frame = self.cap.read()
            if ret:
                current_frame = frame # 전송을 위해 프레임 저장
                # [분류 실행]
                detected_label = run_detection_once(frame)
        if detected_label == 'None':
            detected_label = 'general'
        # ---------------------------------------------------------
        # 백엔드 전송 (분류 직후 실행)
        # ---------------------------------------------------------
        if current_frame is not None:
            # UI 버벅임 방지를 위해 스레드로 실행
            t = threading.Thread(
                target=process_and_send_backend, 
                args=(detected_label, current_frame)
            )
            t.daemon = True # 앱 종료시 스레드도 자동 종료
            t.start()
        # ---------------------------------------------------------

        # [STEP 3] 분류 결과에 따른 서브모터 명령
        servo_command = 'G'
        if detected_label == 'can':
            servo_command = 'C'
        elif detected_label == 'plastic cup': 
            servo_command = 'P'
            
        print(f"1. 분류 결과: {detected_label} -> 서브모터 명령: {servo_command}")
        
        motor.send_command(servo_command)

        
        self.last_detected_label = detected_label
        self.subtitle_label.setText(f"결과: {detected_label}\n서브모터 이동 중...")
        
        # [STEP 4] 2초 후 압축 실행
        QTimer.singleShot(2000, self.run_compression_step)

    def run_compression_step(self):
        """[STEP 4] 압축 모터 작동"""
        print("2. 압축 모터(메인) 작동 시작")
        
        # 메인 모터 명령 (필요 시 주석 해제)
        # motor.send_command('Z') 
        
        self.subtitle_label.setText("압축 진행 중...")
        
        self.timer.start(30)
        QTimer.singleShot(10000, self.stop_process)

    def update_frame(self):
        if self.cap and self.cap.isOpened():
            ret, frame = self.cap.read()
            if ret:
                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                frame = cv2.resize(frame, (320, 240))
                h, w, c = frame.shape
                q_img = QImage(frame.data, w, h, c * w, QImage.Format.Format_RGB888)
                self.video_label.setPixmap(QPixmap.fromImage(q_img))

    def stop_process(self):
        self.timer.stop()
        self.video_label.hide()
        self.top_btn.show()
        label = getattr(self, 'last_detected_label', 'Unknown')
        self.subtitle_label.setText(f"처리 완료! (이전: {label})")
        self.is_processing = False
        self.update_waste_status()

    def update_waste_status(self):
        data = sensor.get_data() 
        if data:
            gen_val = int(data.get("General", 0) * 100)
            can_val = int(data.get("Can", 0) * 100)
            pla_val = int(data.get("Plastic", 0) * 100)
            self.general_row.percent.setText(f"{gen_val}%")
            self.can_row.percent.setText(f"{can_val}%")
            self.plastic_row.percent.setText(f"{pla_val}%")

    def closeEvent(self, event):
        if self.timer.isActive():
            self.timer.stop()
        if self.cap and self.cap.isOpened():
            self.cap.release()
        event.accept()

if __name__ == '__main__':
    app = QApplication(sys.argv)
    ex = WasteApp()
    ex.show()
    sys.exit(app.exec())