import streamlit as st
import cv2
from ultralytics import YOLO
from utils.counter import Counter
from utils.zones import draw_count_line_horizontal, draw_count_line_vertical
import tempfile
import time

from metrics import get_metrics

metrics = get_metrics()

st.set_page_config(layout="wide", page_title="People Counter")

STREAM_URL = "http://192.168.11.76:8080?action=stream"

start_button = st.button("Start Counter")

FRAME_SKIP = 2  # Reduziert Verarbeitungslast

if start_button and STREAM_URL:
    line_orientation = 'vertical'
    line_position = 640

    stframe = st.empty()
    model = YOLO("yolo_weights/yolo11n.pt")
    counter = Counter(line_orientation, metrics)

    cap = cv2.VideoCapture(STREAM_URL)

    frame_idx = 0
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            st.warning("Kein Bild vom RTSP-Stream.")
            break

        frame_idx += 1
        if frame_idx % FRAME_SKIP != 0:
            continue

        results = model.track(frame, persist=True, classes=0, tracker="bytetrack.yaml")
        boxes = results[0].boxes

        if boxes.id is not None:
            for box, track_id in zip(boxes.xywh.cpu(), boxes.id.cpu()):
                x, y, w, h = map(int, box)
                cv2.rectangle(frame, (x-w//2, y-h//2), (x+w//2, y+h//2), (0, 255, 0), 2)
                cv2.putText(frame, str(track_id), (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0,255,0), 2)

            tracks = [
                type("Track", (), {"id": int(track_id), "bbox": box.numpy()})
                for box, track_id in zip(boxes.xywh.cpu(), boxes.id.cpu())
            ]
            counter.update(tracks, line_position)

        if line_orientation == 'horizontal':
            draw_count_line_horizontal(frame, line_position)
        elif line_orientation == 'vertical':
            draw_count_line_vertical(frame, line_position)
        else:
            raise NotImplementedError(f'Line orientation {line_orientation} is invalid!')
        cv2.putText(frame, f"In: {counter.in_count}", (10, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 2)
        cv2.putText(frame, f"Out: {counter.out_count}", (10, 80), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 2)

        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        stframe.image(frame, channels="RGB")

    cap.release()
