import cv2
import mediapipe as mp
import numpy as np
import streamlit as st
from PIL import Image

# Streamlit Setup
st.set_page_config(page_title="AI Virtual Painter", layout="wide")
st.title("AI Virtual Painter")
st.sidebar.header("Settings")

# Sidebar Settings
draw_color = st.sidebar.color_picker("Select Drawing Color", "#FF0000")
thickness = st.sidebar.slider("Brush Thickness", 5, 50, 20)
clear_canvas = st.sidebar.button("Clear Canvas")

# Mediapipe Setup
mp_hands = mp.solutions.hands
hands = mp_hands.Hands(min_detection_confidence=0.85, min_tracking_confidence=0.5, max_num_hands=1)

# Canvas Dimensions
width, height = 1280, 720
canvas = np.zeros((height, width, 3), np.uint8)
current_color = (255, 0, 0)
xp, yp = 0, 0  # Previous points

# Clear Canvas Logic
if clear_canvas:
    canvas = np.zeros((height, width, 3), np.uint8)

# Webcam Feed
cap = cv2.VideoCapture(0)  # Capture from webcam (camera index 0)

# Run Webcam
if cap.isOpened():
    st.text("Press 'q' in the webcam feed window to quit.")

    while True:
        ret, frame = cap.read()
        if not ret:
            st.error("Failed to grab frame from webcam.")
            break

        # Flip and Convert Frame
        frame = cv2.flip(frame, 1)
        rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

        # Process Frame for Hand Landmarks
        results = hands.process(rgb_frame)
        if results.multi_hand_landmarks:
            for hand_landmarks in results.multi_hand_landmarks:
                for id, lm in enumerate(hand_landmarks.landmark):
                    cx, cy = int(lm.x * width), int(lm.y * height)

                    # Tip of Index Finger
                    if id == 8:
                        if xp == 0 and yp == 0:
                            xp, yp = cx, cy

                        # Draw on Canvas
                        cv2.line(canvas, (xp, yp), (cx, cy), current_color, thickness)
                        xp, yp = cx, cy

        # Combine Canvas and Frame
        combined_frame = cv2.addWeighted(frame, 0.5, canvas, 0.5, 0)

        # Display Frame with OpenCV
        cv2.imshow("AI Virtual Painter", combined_frame)

        # Break Loop with 'q'
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

# Cleanup
cap.release()
cv2.destroyAllWindows()
hands.close()