import cv2
import time
import json
import subprocess
import threading
import base64
import requests
from pathlib import Path
import RPi.GPIO as GPIO
from openai import OpenAI
import numpy as np
from lerobot.common.robot_devices.motors.feetech import FeetechMotorsBus, TorqueMode
from lerobot.common.robot_devices.motors.configs import FeetechMotorsBusConfig
import json
from pathlib import Path



ROBOT_SEQUENCE = "throw_candy_at_kids" # this is your sequence name that you want to use
ROBOT_SPEED = "slow"  # very_slow, slow, medium, fast
ROBOT_PORT = "/dev/ttyACM0"  # Change based on your Pi setup

# Ultrasonic sensor pins
TRIG_PIN = 18
ECHO_PIN = 24

# Detection settings
DETECTION_DISTANCE_FEET = 10
DETECTION_DISTANCE_CM = DETECTION_DISTANCE_FEET * 30.48  # Convert to cm

# Camera settings
CAMERA_INDEX = 0  # Usually 0 for USB webcam

# OpenAI API settings (set your API key)
OPENAI_API_KEY = "your-openai-api-key-here"  # CHANGE THIS!

# Text-to-speech settings
USE_ESPEAK = True  # Set to False to use festival instead

# =============================================================================
# CANDY DISPENSER CLASS
# =============================================================================

class CandyDispenser:
    def __init__(self):
        self.client = OpenAI(api_key=OPENAI_API_KEY)
        self.camera = None
        self.setup_gpio()
        self.setup_camera()
        self.robot_busy = False
        
    def setup_gpio(self):
        """Setup GPIO pins for ultrasonic sensor"""
        GPIO.setmode(GPIO.BCM)
        GPIO.setup(TRIG_PIN, GPIO.OUT)
        GPIO.setup(ECHO_PIN, GPIO.IN)
        print(f"GPIO setup complete - Ultrasonic sensor on pins {TRIG_PIN}/{ECHO_PIN}")
    
    def setup_camera(self):
        """Setup webcam"""
        try:
            self.camera = cv2.VideoCapture(CAMERA_INDEX)
            if not self.camera.isOpened():
                raise Exception("Camera not found")
            print(f"Camera setup complete on index {CAMERA_INDEX}")
        except Exception as e:
            print(f"Camera setup failed: {e}")
            self.camera = None
    
    def measure_distance(self):
        """Measure distance using ultrasonic sensor"""
        try:
            # Send trigger pulse
            GPIO.output(TRIG_PIN, True)
            time.sleep(0.00001)  # 10 microseconds
            GPIO.output(TRIG_PIN, False)
            
            # Measure echo time
            start_time = time.time()
            while GPIO.input(ECHO_PIN) == 0:
                pulse_start = time.time()
                if pulse_start - start_time > 0.1:  # Timeout after 100ms
                    return float('inf')
            
            while GPIO.input(ECHO_PIN) == 1:
                pulse_end = time.time()
                if pulse_end - pulse_start > 0.1:  # Timeout after 100ms
                    return float('inf')
            
            # Calculate distance
            pulse_duration = pulse_end - pulse_start
            distance_cm = pulse_duration * 17150  # Speed of sound calculation
            
            return distance_cm
            
        except Exception as e:
            print(f"Distance measurement error: {e}")
            return float('inf')
            
    def play_sequence_direct(
        self,
        name=ROBOT_SEQUENCE,
        speed_mode=ROBOT_SPEED,
        port=ROBOT_PORT,
        sequence_folder="sequences"
    ):
        """Directly play a LeRobot sequence using internal code, no subprocess or input()"""
        # This is nearly a copy of your playback code, but refactored to be self-contained.
        motors_config = {
            "shoulder_pan": [1, "sts3215"],
            "shoulder_lift": [2, "sts3215"],
            "elbow_flex": [3, "sts3215"],
            "wrist_flex": [4, "sts3215"],
            "wrist_roll": [5, "sts3215"],
            "gripper": [6, "sts3215"],
        }

        file_path = Path(sequence_folder) / f"{name}.json"
        if not file_path.exists():
            print(f"Sequence file not found: {file_path}")
            return False

        with open(file_path, "r") as f:
            data = json.load(f)
        sequence = data["sequence"]

        # Connect to motors
        config = FeetechMotorsBusConfig(port=port, motors=motors_config)
        motor_bus = FeetechMotorsBus(config)
        motor_bus.connect()

        # REMOVE ALL LIMITS
        print("REMOVING ALL LIMITS...")
        for motor_name in motors_config.keys():
            try:
                motor_bus.write("Min_Angle_Limit", 0, motor_name)
                motor_bus.write("Max_Angle_Limit", 4095, motor_name)
                print(f"   {motor_name}: Limits removed")
            except Exception as e:
                print(f"   {motor_name}: {e}")
        print("Connected and ALL limits removed")

        # Set speed mode (same as your script)
        speed_settings = {
            "very_slow": {"acceleration": 50, "max_accel": 50, "description": "Very Slow (50)"},
            "slow": {"acceleration": 100, "max_accel": 100, "description": "Slow (100)"},
            "medium": {"acceleration": 150, "max_accel": 150, "description": "Medium (150)"},
            "fast": {"acceleration": 254, "max_accel": 254, "description": "Fast (254 - Official LeRobot)"},
        }
        if speed_mode not in speed_settings:
            speed_mode = "slow"
        settings = speed_settings[speed_mode]

        print(f"\nSETTING SPEED: {settings['description']}")
        for motor_name in motors_config.keys():
            try:
                motor_bus.write("Mode", 0, motor_name)  # Position Control
                motor_bus.write("P_Coefficient", 16, motor_name)  # Smooth movement
                motor_bus.write("I_Coefficient", 0, motor_name)
                motor_bus.write("D_Coefficient", 32, motor_name)
                motor_bus.write("Lock", 0, motor_name)  # Unlock EPROM
                motor_bus.write("Maximum_Acceleration", settings["max_accel"], motor_name)
                motor_bus.write("Acceleration", settings["acceleration"], motor_name)
                print(f"   {motor_name}: Speed set to {settings['description']}")
            except Exception as e:
                print(f"   {motor_name}: {e}")

        # Torque ON
        motor_bus.write("Torque_Enable", TorqueMode.ENABLED.value)
        print("Torque ON - robot under control")

        # Run sequence!
        try:
            print(f"\nPLAYING SEQUENCE: {name}")
            for i, step in enumerate(sequence):
                pos = step["positions"]
                duration = step["duration"]
                position_num = step["position"]
                pos_str = " | ".join([f"{motor}:{position:4d}" for motor, position in pos.items()])
                if position_num == 1:
                    print(f"Position {position_num}: {pos_str} (moving to start)")
                    for motor_name, position in pos.items():
                        motor_bus.write("Goal_Position", int(position), motor_name)
                    time.sleep(1.5)
                else:
                    if duration == 0:
                        print(f"Position {position_num}: {pos_str} (SUPER FAST)")
                        for motor_name, position in pos.items():
                            motor_bus.write("Goal_Time", 200, motor_name)
                            motor_bus.write("Goal_Position", int(position), motor_name)
                        time.sleep(0.4)
                    else:
                        print(f"Position {position_num}: {pos_str} (SMOOTH ~{duration}s)")
                        for motor_name, position in pos.items():
                            motor_bus.write("Goal_Position", int(position), motor_name)
                        time.sleep(0.8)
            print(f"\nPROPER SLOW SEQUENCE COMPLETE!")
            return True
        except Exception as e:
            print(f"ERROR during playback: {e}")
            return False
        finally:
            try:
                motor_bus.disconnect()
                print("Disconnected")
            except Exception as e:
                print(f"Error on disconnect: {e}")
    
    def capture_image(self, delay_seconds=2):
        """Capture an image from the webcam, save, and return base64 string."""
        if not self.camera or not self.camera.isOpened():
            raise RuntimeError("Could not open camera")

        self.camera.set(cv2.CAP_PROP_AUTO_EXPOSURE, 0.75)
        self.camera.read()  # Dummy read to trigger auto exposure

        time.sleep(delay_seconds)

        for _ in range(5):  # Flush buffer
            self.camera.read()

        ret, frame = self.camera.read()

        if not ret:
            raise RuntimeError("Failed to capture image from camera")

        # Brightness boost (optional)
        hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
        h, s, v = cv2.split(hsv)
        v = cv2.add(v, 50)
        v = np.clip(v, 0, 255)
        final_hsv = cv2.merge((h, s, v))
        frame = cv2.cvtColor(final_hsv, cv2.COLOR_HSV2BGR)

        # Save image
        timestamp = time.strftime("%Y%m%d-%H%M%S")
        image_path = f"/home/donuts/lerobot/captures/visitor_{timestamp}.jpg"
        Path("/home/donuts/lerobot/captures").mkdir(parents=True, exist_ok=True)
        cv2.imwrite(str(image_path), frame)

        # Encode to base64
        retval, buffer = cv2.imencode('.jpg', frame)
        if not retval:
            raise RuntimeError("Failed to encode image")

        encoded_image = base64.b64encode(buffer).decode('utf-8')
        return encoded_image

    def analyze_with_chatgpt(self, image_base64):
        """Send image to OpenAI's Vision model using the new responses endpoint"""
        try:
            response = self.client.responses.create(
                model="gpt-4o",  # or "gpt-4.1" if you have access
                input=[{
                    "role": "user",
                    "content": [
                        {"type": "input_text", "text": """Analyze this image for a candy dispenser robot. 

RESPOND IN EXACTLY THIS FORMAT:
ACTION: [CANDY or NO_CANDY]
RESPONSE: [Your brief, funny response]

Rules:
- If you see a child/kid (anyone who looks under 16), respond with ACTION: CANDY
- If you see an adult or no person, respond with ACTION: NO_CANDY
- Keep RESPONSE to 1-2 sentences maximum
- Be funny if possible
- For kids: comment on their costume if applicable or just comment about something you see like attire or appearance (dont be mean)
- For adults: playfully deny them candy in a funny way
- If there are kids and adults, comment on the kids and do give them candy
- If there is no one there, you cant identify the image, or something of that sort, give no candy and comment according to what you see as best you can in the same fun way

Examples:
ACTION: CANDY
RESPONSE: Hey there little vampire! Nice cape! Here's some candy to fuel your spooky adventures!

ACTION: NO_CANDY  
RESPONSE: Sorry grown-up, no sugar rush for you today!

Be sure to stick to that format, but be creative when possible with your response, and comment on things that make it clear you're actually seeing who you're looking at.""" },
                        {
                            "type": "input_image",
                            "image_url": f"data:image/jpeg;base64,{image_base64}",
                            "detail": "high"
                        }
                    ],
                }]
            )

            return response.output_text.strip()

        except Exception as e:
            print(f"ChatGPT analysis error: {e}")
            return "ACTION: NO_CANDY\nRESPONSE: Oops! My robot brain is having a glitch. Try again in a moment!"

    def parse_chatgpt_response(self, response):
        """Parse ChatGPT response into action and speech"""
        try:
            lines = response.split('\n')
            action_line = None
            response_line = None
            
            for line in lines:
                if line.startswith('ACTION:'):
                    action_line = line.split('ACTION:')[1].strip()
                elif line.startswith('RESPONSE:'):
                    response_line = line.split('RESPONSE:')[1].strip()
            
            if not action_line or not response_line:
                print(f"Failed to parse response: {response}")
                return "NO_CANDY", "Sorry, I'm having trouble thinking right now!"
            
            should_give_candy = action_line.upper() == "CANDY"
            return "CANDY" if should_give_candy else "NO_CANDY", response_line
            
        except Exception as e:
            print(f"Response parsing error: {e}")
            return "NO_CANDY", "My circuits are a bit scrambled right now!"
    
    def speak_text(self, text):
        """Use OpenAI's TTS to generate high-quality voice and play it"""
        try:
            response = self.client.audio.speech.create(
                model="tts-1",  # Or use "tts-1-hd" for higher fidelity
                voice="echo",   # Other options: alloy, echo, fable, nova, shimmer, onyx
                input=text
            )

            audio_path = "/tmp/response.mp3"
            with open(audio_path, "wb") as f:
                f.write(response.content)

            subprocess.run(["mpg123", audio_path], check=True)
            print(f"Spoke: {text}")
        except Exception as e:
            print(f"TTS error: {e}")
            print(f"Would have said: {text}")

    def handle_visitor(self):
        """Handle a detected visitor"""
        print("Visitor detected! Analyzing...")
        
        image_base64 = self.capture_image()

        # Analyze with ChatGPT
        print("Analyzing image with ChatGPT...")
        chatgpt_response = self.analyze_with_chatgpt(image_base64)
        print(f"ChatGPT response: {chatgpt_response}")
        
        # Parse response
        action, speech_text = self.parse_chatgpt_response(chatgpt_response)
        
        # Speak the response
        self.speak_text(speech_text)
        
        # Dispense candy if appropriate
        if action == "CANDY":
            print("Giving candy!")
            success = self.play_sequence_direct()

            if not success:
                self.speak_text("Oops! My arm got stuck. The candy is stuck too, but you deserve it anyway!")
        else:
            print("No candy this time")
        
    
    def run(self):
        """Main loop"""
        print(f"\nHALLOWEEN CANDY DISPENSER ACTIVE")
        print(f"Sequence: {ROBOT_SEQUENCE} | Speed: {ROBOT_SPEED} | Port: {ROBOT_PORT}")
        print(f"Detection range: {DETECTION_DISTANCE_FEET} feet")
        print(f"Watching for trick-or-treaters...")
        print("Press Ctrl+C to stop\n")
        
        last_detection_time = 0
        cooldown_period = 50  # seconds between detections
        
        try:
            while True:
                distance_cm = self.measure_distance()
                current_time = time.time()
                
                if distance_cm <= DETECTION_DISTANCE_CM:
                    if current_time - last_detection_time > cooldown_period:
                        print(f"Person detected at {distance_cm:.1f}cm!")
                        last_detection_time = current_time
                        
                        # Handle visitor in separate thread to avoid blocking
                        visitor_thread = threading.Thread(target=self.handle_visitor)
                        visitor_thread.start()
                
                time.sleep(0.1)  # Check 10 times per second
                
        except KeyboardInterrupt:
            print("\nShutting down candy dispenser...")
        finally:
            self.cleanup()
    
    def cleanup(self):
        """Clean up resources"""
        try:
            if self.camera:
                self.camera.release()
            GPIO.cleanup()
            print("Cleanup complete")
        except Exception as e:
            print(f"Cleanup error: {e}")

# =============================================================================
# MAIN EXECUTION
# =============================================================================

if __name__ == "__main__":
    # Check if OpenAI API key is set
    if OPENAI_API_KEY == "your-openai-api-key-here":
        print("ERROR: Please set your OpenAI API key in the OPENAI_API_KEY variable!")
        exit(1)
    
    # Create and run candy dispenser
    dispenser = CandyDispenser()
    dispenser.run()