#!/usr/bin/env python3
"""
Retail Video Analysis Pipeline
Analyze retail store videos to extract product prices and inventory information.

Usage:
    python analyze_retail_video.py <video_path> [--output <dir>] [--start <time>] [--end <time>]
    [--use-gemini] [--lang <languages>] [--debug] [--save-frames]

Example:
    python analyze_retail_video.py store_video.mp4 --output ./analysis --use-gemini
"""

import os
import sys
import argparse
import subprocess
import json
import re
from datetime import datetime
from pathlib import Path

# Try to import optional dependencies
try:
    from PIL import Image, ImageEnhance, ImageFilter
    HAS_PIL = True
except ImportError:
    HAS_PIL = False
    print("Warning: Pillow not installed. Install with: pip install Pillow")

try:
    import pytesseract
    HAS_TESSERACT = True
except ImportError:
    HAS_TESSERACT = False
    print("Warning: pytesseract not installed. Install with: pip install pytesseract")

try:
    import google.generativeai as genai
    HAS_GEMINI = True
except ImportError:
    HAS_GEMINI = False
    print("Warning: google-generativeai not installed. Install with: pip install google-generativeai")

# Price pattern regexes for different formats
PRICE_PATTERNS = [
    r'€\s*\d+[.,]\d{2}',      # €12.99, € 12,99
    r'\$\s*\d+[.,]\d{2}',     # $12.99, $ 12.99
    r'\d+[.,]\d{2}\s*€',      # 12.99 €, 12,99€
    r'\d+[.,]\d{2}\s*\$',     # 12.99 $, 12,99$
    r'\d+\s*[Ee]uros?',       # 12 Euros, 12 euros
    r'\d+\s*[Dd]ollars?',     # 12 Dollars, 12 dollars
]

def extract_frames(video_path, output_dir, start_time=None, end_time=None):
    """
    Extract frames from video using ffmpeg.
    Returns list of extracted frame paths.
    """
    frames = []
    
    # Create output directory if it doesn't exist
    os.makedirs(output_dir, exist_ok=True)
    
    # Get video duration
    cmd = ['ffprobe', '-v', 'error', '-show_entries', 'format=duration',
           '-of', 'default=noprint_wrappers=1:nokey=1', video_path]
    try:
        duration = float(subprocess.check_output(cmd).decode().strip())
    except:
        duration = 60  # Default to 60 seconds if can't determine
    
    # Determine frame extraction times
    if start_time and end_time:
        # Parse time strings (HH:MM:SS or MM:SS)
        def parse_time(t):
            parts = list(map(int, t.split(':')))
            if len(parts) == 2:  # MM:SS
                return parts[0] * 60 + parts[1]
            elif len(parts) == 3:  # HH:MM:SS
                return parts[0] * 3600 + parts[1] * 60 + parts[2]
            return 0
        
        start_sec = parse_time(start_time)
        end_sec = parse_time(end_time)
        extraction_times = [start_sec, (start_sec + end_sec) / 2, end_sec]
    else:
        # Extract at start, middle, and end
        extraction_times = [0, duration / 2, duration - 1]
    
    # Extract frames
    for i, time_sec in enumerate(extraction_times):
        if time_sec < 0 or time_sec > duration:
            continue
            
        frame_path = os.path.join(output_dir, f'frame_{i:02d}.jpg')
        cmd = [
            'ffmpeg', '-ss', str(time_sec), '-i', video_path,
            '-frames:v', '1', '-q:v', '2', '-y', frame_path
        ]
        
        try:
            subprocess.run(cmd, check=True, capture_output=True)
            frames.append({
                'path': frame_path,
                'timestamp': time_sec,
                'time_str': f'{int(time_sec//3600):02d}:{int((time_sec%3600)//60):02d}:{int(time_sec%60):02d}'
            })
            print(f"Extracted frame: {frame_path} (at {frames[-1]['time_str']})")
        except subprocess.CalledProcessError as e:
            print(f"Error extracting frame at {time_sec}s: {e}")
    
    return frames

def preprocess_image(image_path):
    """
    Preprocess image for better OCR results.
    """
    if not HAS_PIL:
        return image_path
    
    try:
        img = Image.open(image_path)
        
        # Convert to grayscale
        img = img.convert('L')
        
        # Enhance contrast
        enhancer = ImageEnhance.Contrast(img)
        img = enhancer.enhance(2.0)
        
        # Sharpen slightly
        img = img.filter(ImageFilter.SHARPEN)
        
        # Save processed image
        processed_path = image_path.replace('.jpg', '_processed.jpg')
        img.save(processed_path, 'JPEG', quality=95)
        
        return processed_path
    except Exception as e:
        print(f"Error preprocessing image {image_path}: {e}")
        return image_path

def extract_text_with_ocr(image_path, languages='eng+grc'):
    """
    Extract text from image using Tesseract OCR.
    """
    if not HAS_TESSERACT:
        return ""
    
    try:
        # Preprocess image first
        processed_path = preprocess_image(image_path)
        
        # Extract text
        text = pytesseract.image_to_string(
            Image.open(processed_path),
            lang=languages,
            config='--psm 6 --oem 3'
        )
        
        return text.strip()
    except Exception as e:
        print(f"Error in OCR for {image_path}: {e}")
        return ""

def find_prices_in_text(text):
    """
    Find price patterns in extracted text.
    """
    prices = []
    
    for pattern in PRICE_PATTERNS:
        matches = re.finditer(pattern, text, re.IGNORECASE)
        for match in matches:
            price_text = match.group(0)
            # Clean up the price text
            price_text = re.sub(r'\s+', '', price_text)
            prices.append({
                'price': price_text,
                'raw_match': match.group(0),
                'pattern': pattern
            })
    
    return prices

def analyze_with_gemini(image_path, api_key=None):
    """
    Use Gemini Vision for enhanced analysis.
    """
    if not HAS_GEMINI:
        return {"error": "Gemini not available"}
    
    try:
        # Configure Gemini
        if api_key:
            genai.configure(api_key=api_key)
        else:
            # Try to get from environment
            api_key = os.getenv('GOOGLE_API_KEY')
            if not api_key:
                return {"error": "No Gemini API key provided"}
            genai.configure(api_key=api_key)
        
        # Load and analyze image
        model = genai.GenerativeModel('gemini-2.0-flash-exp')
        
        with open(image_path, 'rb') as f:
            image_data = f.read()
        
        prompt = """Analyze this retail store image and provide:
        1. Product categories visible (e.g., wine, bread, dairy)
        2. Price tags and their values
        3. Store organization and shelf layout
        4. Any promotional displays or signs
        5. Overall store condition and presentation
        
        Format the response as a JSON object with these keys:
        - product_categories: array of strings
        - price_tags: array of objects with 'price' and 'product' fields
        - store_organization: string description
        - promotional_displays: array of strings
        - store_condition: string (poor/fair/good/excellent)
        """
        
        response = model.generate_content([
            prompt,
            {"mime_type": "image/jpeg", "data": image_data}
        ])
        
        # Try to parse JSON from response
        try:
            # Extract JSON from response text
            text = response.text
            # Find JSON object in text
            json_start = text.find('{')
            json_end = text.rfind('}') + 1
            if json_start != -1 and json_end > json_start:
                json_str = text[json_start:json_end]
                return json.loads(json_str)
            else:
                return {"analysis": text}
        except:
            return {"analysis": response.text}
            
    except Exception as e:
        return {"error": str(e)}

def analyze_video(video_path, args):
    """
    Main analysis pipeline.
    """
    print(f"Analyzing video: {video_path}")
    
    # Setup output directory
    if args.output:
        output_dir = args.output
    else:
        video_name = Path(video_path).stem
        output_dir = f"./retail_analysis_{video_name}_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
    
    os.makedirs(output_dir, exist_ok=True)
    
    # Extract frames
    print("Step 1: Extracting frames from video...")
    frames = extract_frames(video_path, output_dir, args.start, args.end)
    
    if not frames:
        print("Error: No frames extracted")
        return None
    
    # Analyze each frame
    all_results = {
        'video_info': {
            'path': video_path,
            'analysis_timestamp': datetime.now().isoformat(),
            'frames_analyzed': len(frames)
        },
        'frame_analyses': [],
        'summary': {
            'total_prices_found': 0,
            'price_range': None,
            'product_categories': set(),
            'gemini_analyses': 0
        }
    }
    
    prices_found = []
    
    for i, frame in enumerate(frames):
        print(f"\nStep 2: Analyzing frame {i+1}/{len(frames)} ({frame['time_str']})...")
        
        frame_result = {
            'frame_file': frame['path'],
            'timestamp': frame['time_str'],
            'ocr_text': '',
            'prices': [],
            'gemini_analysis': None
        }
        
        # OCR analysis
        if HAS_TESSERACT:
            print("  - Running OCR...")
            ocr_text = extract_text_with_ocr(frame['path'], args.lang)
            frame_result['ocr_text'] = ocr_text
            
            # Find prices in OCR text
            prices = find_prices_in_text(ocr_text)
            frame_result['prices'] = prices
            prices_found.extend(prices)
            
            if prices:
                print(f"  - Found {len(prices)} price(s): {[p['price'] for p in prices]}")
            else:
                print("  - No prices found in OCR text")
        
        # Gemini Vision analysis (if requested)
        if args.use_gemini and HAS_GEMINI:
            print("  - Running Gemini Vision analysis...")
            gemini_result = analyze_with_gemini(frame['path'], args.gemini_key)
            frame_result['gemini_analysis'] = gemini_result
            
            if 'error' not in gemini_result:
                all_results['summary']['gemini_analyses'] += 1
                print("  - Gemini analysis completed")
        
        all_results['frame_analyses'].append(frame_result)
    
    # Generate summary
    if prices_found:
        # Extract numeric values for range calculation
        numeric_prices = []
        for price in prices_found:
            # Extract numbers from price string
            nums = re.findall(r'\d+[.,]\d+|\d+', price['price'])
            if nums:
                # Convert to float (handle comma as decimal separator)
                num_str = nums[0].replace(',', '.')
                try:
                    numeric_prices.append(float(num_str))
                except:
                    pass
        
        if numeric_prices:
            min_price = min(numeric_prices)
            max_price = max(numeric_prices)
            all_results['summary']['price_range'] = f"{min_price:.2f} - {max_price:.2f}"
    
    all_results['summary']['total_prices_found'] = len(prices_found)
    
    # Save results
    results_file = os.path.join(output_dir, 'analysis_results.json')
    with open(results_file, 'w', encoding='utf-8') as f:
        json.dump(all_results, f, indent=2, ensure_ascii=False)
    
    # Generate human-readable report
    report_file = os.path.join(output_dir, 'analysis_report.txt')
    with open(report_file, 'w', encoding='utf-8') as f:
        f.write(f"RETAIL VIDEO ANALYSIS REPORT\n")
        f.write(f"=============================\n")
        f.write(f"Video: {video_path}\n")
        f.write(f"Analysis Date: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
        f.write(f"Frames Analyzed: {len(frames)}\n")
        f.write(f"\nSUMMARY\n")
        f.write(f"-------\n")
        f.write(f"Total Prices Found: {len(prices_found)}\n")
        if all_results['summary']['price_range']:
            f.write(f"Price Range: {all_results['summary']['price_range']}\n")
        
        f.write(f"\nDETAILED FINDINGS\n")
        f.write(f"-----------------\n")
        
        for i, frame_analysis in enumerate(all_results['frame_analyses']):
            f.write(f"\nFrame {i+1} ({frame_analysis['timestamp']}):\n")
            f.write(f"  Prices found: {len(frame_analysis['prices'])}\n")
            for price in frame_analysis['prices']:
                f.write(f"    - {price['price']}\n")
            
            if frame_analysis['gemini_analysis'] and 'error' not in frame_analysis['gemini_analysis']:
                f.write(f"  Gemini Analysis:\n")
                gemini = frame_analysis['gemini_analysis']
                if 'product_categories' in gemini:
                    f.write(f"    Product Categories: {', '.join(gemini['product_categories'])}\n")
                if 'price_tags' in gemini:
                    f.write(f"    Price Tags: {len(gemini['price_tags'])} found\n")
    
    print(f"\nAnalysis complete!")
    print(f"Results saved to: {output_dir}/")
    print(f"  - JSON results: {results_file}")
    print(f"  - Human report: {report_file}")
    
    return all_results

def main():
    parser = argparse.ArgumentParser(description='Analyze retail store videos for price and inventory information')
    parser.add_argument('video_path', help='Path to the video file')
    parser.add_argument('--output', '-o', help='Output directory for analysis results')
    parser.add_argument('--start', '-s', help='Start time for analysis (HH:MM:SS or MM:SS)')
    parser.add_argument('--end', '-e', help='End time for analysis (HH:MM:SS or MM:SS)')
    parser.add_argument('--use-gemini', '-g', action='store_true', help='Use Gemini Vision for enhanced analysis')
    parser.add_argument('--gemini-key', help='Gemini API key (or set GOOGLE_API_KEY env var)')
    parser.add_argument('--lang', default='eng+grc', help='OCR languages (default: eng+grc for English+Greek)')
    parser.add_argument('--debug', action='store_true', help='Enable debug output')
    parser.add_argument('--save-frames', action='store_true', help='Save preprocessed frames')
    
    args = parser.parse_args()
    
    # Check if video file exists
    if not os.path.exists(args.video_path):
        print(f"Error: Video file not found: {args.video_path}")
        sys.exit(1)
    
    # Run analysis
    results = analyze_video(args.video_path, args)
    
    if results:
        # Print quick summary
        print(f"\nQUICK SUMMARY:")
        print(f"  Total frames analyzed: {results['video_info']['frames_analyzed']}")
        print(f"  Total prices found: {results['summary']['total_prices_found']}")
        if results['summary']['price_range']:
            print(f"  Price range: {results['summary']['price_range']}")
        
        if args.use_gemini:
            print(f"  Gemini analyses: {results['summary']['gemini_analyses']}")
    else:
        print("Analysis failed")
        sys.exit(1)

if __name__ == '__main__':
    main()