#!/usr/bin/env python3
"""
CFP-Jarvis1 Streaming Mode - Real-time screen analysis and action
"""

import os
import time
import threading
import queue
import torch
import numpy as np
from PIL import Image
import mss
import cv2
from transformers import AutoProcessor, AutoModelForImageTextToText
from dataclasses import dataclass
from typing import Optional, Tuple
import pyautogui

@dataclass
class Action:
    """Represents an action to perform"""
    type: str  # 'click', 'type', 'key', 'move'
    coordinates: Optional[Tuple[int, int]] = None
    text: Optional[str] = None
    confidence: float = 0.0

class CFPJarvis1Stream:
    """
    CFP-Jarvis1 with streaming capabilities
    Analyzes screen in real-time and performs actions
    """
    
    def __init__(self, 
                 model_path="OpenGVLab/InternVL3-1B-hf",
                 cache_dir="/media/jerem/641C8D6C1C8D3A56/hf_cache",
                 fps=1,  # Frames per second to analyze
                 action_delay=0.5):  # Delay between actions
        """
        Initialize streaming Jarvis
        
        Args:
            fps: How many frames per second to analyze (1-10 recommended)
            action_delay: Minimum delay between actions in seconds
        """
        self.model_name = "CFP-Jarvis1-Stream"
        self.version = "1.0.0"
        self.fps = min(fps, 10)  # Cap at 10 FPS
        self.action_delay = action_delay
        
        # Streaming components
        self.streaming = False
        self.capture_thread = None
        self.analysis_thread = None
        self.action_thread = None
        
        # Queues for pipeline
        self.frame_queue = queue.Queue(maxsize=10)
        self.action_queue = queue.Queue(maxsize=100)
        
        # Model components
        self.processor = None
        self.model = None
        self.device = None
        self.dtype = None
        
        # Screen capture
        self.sct = mss.mss()
        
        # Performance metrics
        self.last_capture_time = 0
        self.last_analysis_time = 0
        self.frames_analyzed = 0
        
        # Setup cache
        os.environ['HF_HOME'] = cache_dir
        self.model_path = model_path
        
        print(f"🎬 {self.model_name} v{self.version}")
        print(f"📹 Streaming at {self.fps} FPS")
    
    def load_model(self):
        """Load the model for streaming"""
        try:
            print(f"📥 Loading model for streaming...")
            
            self.processor = AutoProcessor.from_pretrained(
                self.model_path,
                trust_remote_code=True
            )
            
            if torch.cuda.is_available():
                self.device = "cuda"
                self.dtype = torch.bfloat16
                print("🚀 GPU streaming enabled")
            else:
                self.device = "cpu"
                self.dtype = torch.float32
                print("💻 CPU streaming (slower)")
            
            self.model = AutoModelForImageTextToText.from_pretrained(
                self.model_path,
                torch_dtype=self.dtype,
                device_map="auto" if self.device == "cuda" else self.device,
                trust_remote_code=True,
                low_cpu_mem_usage=True
            ).eval()
            
            # Compile model for faster inference (PyTorch 2.0+)
            if hasattr(torch, 'compile') and self.device == "cuda":
                print("⚡ Compiling model for faster streaming...")
                self.model = torch.compile(self.model, mode="reduce-overhead")
            
            print("✅ Model ready for streaming")
            return True
            
        except Exception as e:
            print(f"❌ Failed to load model: {e}")
            return False
    
    def capture_screen(self):
        """Continuously capture screen frames"""
        print("📹 Starting screen capture...")
        
        while self.streaming:
            try:
                # Capture screen
                screenshot = self.sct.grab(self.sct.monitors[0])
                
                # Convert to PIL Image
                img = Image.frombytes(
                    'RGB',
                    (screenshot.width, screenshot.height),
                    screenshot.bgra,
                    'raw',
                    'BGRX'
                )
                
                # Resize for faster processing
                img.thumbnail((1280, 720), Image.Resampling.LANCZOS)
                
                # Add to queue if not full
                if not self.frame_queue.full():
                    self.frame_queue.put(img)
                    self.last_capture_time = time.time()
                
                # Control FPS
                time.sleep(1.0 / self.fps)
                
            except Exception as e:
                print(f"❌ Capture error: {e}")
                time.sleep(1)
    
    def analyze_frames(self, task="Monitor screen for actions"):
        """Analyze captured frames and generate actions"""
        print("🔍 Starting frame analysis...")
        
        while self.streaming:
            try:
                # Get frame from queue
                if not self.frame_queue.empty():
                    img = self.frame_queue.get(timeout=1)
                    
                    # Quick analysis prompt for streaming
                    prompt = f"Task: {task}\nWhat is the most important UI element visible? Provide coordinates if clickable."
                    
                    # Generate response
                    start_time = time.time()
                    response = self._quick_analyze(img, prompt)
                    analysis_time = time.time() - start_time
                    
                    self.frames_analyzed += 1
                    self.last_analysis_time = analysis_time
                    
                    # Parse response for actions
                    action = self._parse_action(response)
                    
                    if action:
                        self.action_queue.put(action)
                    
                    # Show performance
                    if self.frames_analyzed % 10 == 0:
                        print(f"⚡ FPS: {1/analysis_time:.1f} | Frames: {self.frames_analyzed}")
                
                else:
                    time.sleep(0.1)
                    
            except Exception as e:
                print(f"❌ Analysis error: {e}")
                time.sleep(1)
    
    def _quick_analyze(self, image, prompt):
        """Quick analysis optimized for streaming"""
        try:
            messages = [{
                "role": "user",
                "content": [
                    {"type": "image", "image": image},
                    {"type": "text", "text": prompt}
                ]
            }]
            
            text = self.processor.apply_chat_template(
                messages,
                tokenize=False,
                add_generation_prompt=True
            )
            
            inputs = self.processor(
                text=text,
                images=image,
                return_tensors="pt"
            )
            
            inputs = {
                k: v.to(self.device) if torch.is_tensor(v) else v
                for k, v in inputs.items()
            }
            
            if 'pixel_values' in inputs:
                inputs['pixel_values'] = inputs['pixel_values'].to(self.dtype)
            
            # Fast generation with fewer tokens
            with torch.no_grad():
                outputs = self.model.generate(
                    **inputs,
                    max_new_tokens=50,  # Short response for speed
                    do_sample=False,
                    num_beams=1  # Greedy for speed
                )
            
            response = self.processor.decode(
                outputs[0][inputs['input_ids'].shape[1]:],
                skip_special_tokens=True
            )
            
            return response
            
        except Exception as e:
            return f"Error: {str(e)}"
    
    def _parse_action(self, response):
        """Parse response to extract actionable information"""
        # Simple parsing - would need more sophisticated NLP in production
        action = None
        
        response_lower = response.lower()
        
        if "click" in response_lower:
            # Try to extract coordinates
            import re
            coords = re.findall(r'\((\d+),\s*(\d+)\)', response)
            if coords:
                x, y = int(coords[0][0]), int(coords[0][1])
                action = Action(type="click", coordinates=(x, y), confidence=0.8)
        elif "button" in response_lower or "clickable" in response_lower:
            # Default to center if button mentioned
            action = Action(type="click", coordinates=(640, 360), confidence=0.5)
        
        return action
    
    def execute_actions(self):
        """Execute queued actions"""
        print("🎮 Starting action executor...")
        
        while self.streaming:
            try:
                if not self.action_queue.empty():
                    action = self.action_queue.get(timeout=1)
                    
                    if action.type == "click" and action.coordinates:
                        x, y = action.coordinates
                        print(f"🖱️ Clicking at ({x}, {y})")
                        pyautogui.click(x, y)
                    elif action.type == "type" and action.text:
                        print(f"⌨️ Typing: {action.text}")
                        pyautogui.typewrite(action.text)
                    
                    # Delay between actions
                    time.sleep(self.action_delay)
                else:
                    time.sleep(0.1)
                    
            except Exception as e:
                print(f"❌ Action error: {e}")
                time.sleep(1)
    
    def start_streaming(self, task="Monitor and interact with screen", auto_execute=False):
        """
        Start streaming analysis
        
        Args:
            task: Description of what to monitor/do
            auto_execute: Whether to automatically execute detected actions
        """
        if not self.model:
            if not self.load_model():
                return False
        
        print(f"\n🎬 Starting {self.model_name} streaming mode")
        print(f"📋 Task: {task}")
        print(f"🤖 Auto-execute: {auto_execute}")
        
        self.streaming = True
        
        # Start capture thread
        self.capture_thread = threading.Thread(target=self.capture_screen)
        self.capture_thread.start()
        
        # Start analysis thread
        self.analysis_thread = threading.Thread(
            target=self.analyze_frames,
            args=(task,)
        )
        self.analysis_thread.start()
        
        # Start action thread if auto-execute
        if auto_execute:
            self.action_thread = threading.Thread(target=self.execute_actions)
            self.action_thread.start()
        
        print("✅ Streaming started! Press Ctrl+C to stop.")
        
        return True
    
    def stop_streaming(self):
        """Stop streaming analysis"""
        print("\n🛑 Stopping streaming...")
        self.streaming = False
        
        # Wait for threads to finish
        if self.capture_thread:
            self.capture_thread.join(timeout=2)
        if self.analysis_thread:
            self.analysis_thread.join(timeout=2)
        if self.action_thread:
            self.action_thread.join(timeout=2)
        
        print(f"📊 Stats: {self.frames_analyzed} frames analyzed")
        print("✅ Streaming stopped")
    
    def get_status(self):
        """Get streaming status"""
        return {
            "streaming": self.streaming,
            "frames_analyzed": self.frames_analyzed,
            "queue_size": self.frame_queue.qsize(),
            "actions_pending": self.action_queue.qsize(),
            "avg_fps": 1/self.last_analysis_time if self.last_analysis_time > 0 else 0
        }

# Example usage
if __name__ == "__main__":
    import signal
    
    jarvis = CFPJarvis1Stream(fps=2)  # 2 FPS for testing
    
    # Handle Ctrl+C
    def signal_handler(sig, frame):
        print("\n\n⚠️ Interrupt received")
        jarvis.stop_streaming()
        exit(0)
    
    signal.signal(signal.SIGINT, signal_handler)
    
    # Start streaming
    jarvis.start_streaming(
        task="Monitor screen and identify clickable buttons",
        auto_execute=False  # Don't auto-click for safety
    )
    
    # Keep running
    try:
        while jarvis.streaming:
            time.sleep(1)
            status = jarvis.get_status()
            if status["frames_analyzed"] % 10 == 0 and status["frames_analyzed"] > 0:
                print(f"📊 Status: {status}")
    except KeyboardInterrupt:
        pass