import tkinter as tk
from tkinter import filedialog, messagebox
import cv2
import mediapipe as mp
import numpy as np
import os
from PIL import Image, ImageTk

# Initialize MediaPipe Pose
mp_pose = mp.solutions.pose
mp_drawing = mp.solutions.drawing_utils
pose = mp_pose.Pose(static_image_mode=True, min_detection_confidence=0.5)

class PoseApp:
    def __init__(self, root):
        self.root = root
        self.root.title("OpenPose Maker with Preview")
        self.root.geometry("700x850")
        
        # Variables to hold data
        self.current_skeleton = None
        self.suggested_save_path = ""

        # UI Elements
        self.btn_load = tk.Button(root, text="Step 1: Load Photo", command=self.process_image, 
                                  bg="#4285F4", fg="white", font=('Arial', 11, 'bold'), pady=10)
        self.btn_load.pack(pady=10)

        self.preview_label = tk.Label(root, text="[ Image Preview Area ]", bg="gray90")
        self.preview_label.pack(pady=10, fill=tk.BOTH, expand=True)

        self.status_label = tk.Label(root, text="Ready", fg="gray")
        self.status_label.pack()

        self.btn_save = tk.Button(root, text="Step 2: Save This Pose", command=self.save_image, 
                                  bg="#0F9D58", fg="white", font=('Arial', 11, 'bold'), pady=10, state=tk.DISABLED)
        self.btn_save.pack(pady=20)

    def process_image(self):
        file_path = filedialog.askopenfilename(filetypes=[("Image Files", "*.jpg *.jpeg *.png")])
        if not file_path: return

        image = cv2.imread(file_path)
        if image is None: return
        
        # Process Pose
        image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        results = pose.process(image_rgb)
        
        h, w, _ = image.shape
        self.current_skeleton = np.zeros((h, w, 3), dtype=np.uint8)

        if results.pose_landmarks:
            mp_drawing.draw_landmarks(
                self.current_skeleton, results.pose_landmarks, mp_pose.POSE_CONNECTIONS,
                mp_drawing.DrawingSpec(color=(0, 255, 0), thickness=5, circle_radius=4),
                mp_drawing.DrawingSpec(color=(255, 255, 255), thickness=3)
            )

            # Prepare Auto-Save Path
            base_name, _ = os.path.splitext(file_path)
            self.suggested_save_path = f"{base_name}_pose.png"

            # Fix Tiny Preview: Resize for GUI (max 600px)
            preview_rgb = cv2.cvtColor(self.current_skeleton, cv2.COLOR_BGR2RGB)
            pil_img = Image.fromarray(preview_rgb)
            pil_img.thumbnail((600, 600)) # Resizes while keeping aspect ratio
            
            img_tk = ImageTk.PhotoImage(pil_img)
            self.preview_label.config(image=img_tk, text="")
            self.preview_label.image = img_tk
            
            self.btn_save.config(state=tk.NORMAL)
            self.status_label.config(text="Pose generated! Click Save to finish.", fg="blue")
        else:
            messagebox.showwarning("Error", "No person detected in this photo.")

    def save_image(self):
        if self.current_skeleton is not None and self.suggested_save_path:
            cv2.imwrite(self.suggested_save_path, self.current_skeleton)
            messagebox.showinfo("Saved!", f"File saved as:\n{os.path.basename(self.suggested_save_path)}")
            self.btn_save.config(state=tk.DISABLED) # Prevent double-saving
            self.status_label.config(text=f"Done: {os.path.basename(self.suggested_save_path)}", fg="green")

if __name__ == "__main__":
    root = tk.Tk()
    app = PoseApp(root)
    root.mainloop()