"""
High-performance Frame-specific OCR module using Tesserocr for video processing.

This module provides significantly faster OCR functionality using Tesserocr:
- Pre-loaded Tesseract model for memoization (10-50x faster)
- Direct C++ API access instead of CLI wrapper
- Optimized for video frame batch processing
- Memory-efficient frame handling
- Medical text recognition patterns
- Optimized for endoscopy video overlays

Performance Benefits over pytesseract:
- Model is loaded once at initialization
- No subprocess overhead per frame
- Direct memory access to Tesseract engine
- Batch processing optimizations
"""

import logging
import cv2
import numpy as np
import tesserocr
from PIL import Image, ImageEnhance, ImageFilter
from typing import Dict, Any, Tuple, Optional, List
import re
import threading
import time

logger = logging.getLogger(__name__)

class TesseOCRFrameProcessor:
    """
    High-performance OCR processor using Tesserocr for video frames.
    
    Key Features:
    - Pre-loaded Tesseract model for maximum speed
    - Batch processing for video frames
    - Thread-safe operations
    - Automatic memory management
    - German medical text optimization
    """
    
    def __init__(self, language: str = 'deu+eng'):
        """
        Initialize TesseOCR processor with pre-loaded model.
        
        Args:
            language: Tesseract language code (e.g., 'deu+eng' for German+English)
        """
        self.language = language
        self._lock = threading.Lock()
    
        
        # Initialize Tesseract API - this is the key optimization!
        try:
            self.api = tesserocr.PyTessBaseAPI(lang=language, path=tessdata_path)
            logger.info(f"TesseOCR initialized with language: {language}")
            
            # Set optimal OCR parameters for video frames
            self.api.SetPageSegMode(tesserocr.PSM.SINGLE_BLOCK)  # PSM 6 equivalent
            self.api.SetVariable('tessedit_char_whitelist', 
                               '0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ-üöäÜÖÄß.:/ ')
            
            # Optimize for speed vs accuracy
            self.api.SetVariable('tessedit_enable_dict_correction', '1')
            self.api.SetVariable('tessedit_enable_bigram_correction', '1')
            
        except Exception as e:
            logger.error(f"Failed to initialize TesseOCR: {e}")
            self.api = None
            raise
        
        # Frame-specific OCR configurations
        self.frame_config = {
            'dpi': 300,
            'min_confidence': 30,  # Lower threshold for video frames
            'high_quality_dpi': 400,
            'high_quality_min_confidence': 50
        }
        
        # Medical text patterns for enhanced detection
        self.medical_patterns = [
            r'Patient[:\s]*([A-Za-zäöüÄÖÜß\s\-]+)',
            r'Pat[\.:\s]*([A-Za-zäöüÄÖÜß\s\-]+)',
            r'geb[\.:\s]*(\d{1,2}\.\d{1,2}\.\d{2,4})',
            r'Fall[nr]*[\.:\s]*(\d+)',
            r'Datum[:\s]*(\d{1,2}\.\d{1,2}\.\d{2,4})',
            r'Geschlecht[:\s]*(männlich|weiblich|m|w|male|female)',
            r'Geburtsdatum[:\s]*(\d{1,2}\.\d{1,2}\.\d{2,4})',
            r'(\d{1,2}\.\d{1,2}\.\d{2,4})',  # Generic date pattern
            r'(\d{2,})',  # Case numbers
        ]
        
        # Performance metrics
        self.processed_frames = 0
        self.total_processing_time = 0.0
        
    def __del__(self):
        """Clean up Tesseract API on destruction."""
        if hasattr(self, 'api') and self.api:
            try:
                self.api.End()
            except Exception as e:
                logger.debug(f"Error cleaning up Tesseract API: {e}")
    
    def preprocess_frame_for_ocr(self, frame: np.ndarray, roi: Optional[Dict[str, Any]] = None) -> Image.Image:
        """
        Preprocess a video frame for optimal OCR results.
        
        Args:
            frame: Input frame as numpy array (BGR format from cv2)
            roi: Optional region of interest to focus OCR on
            
        Returns:
            Preprocessed PIL Image ready for OCR
        """
        try:
            # Convert BGR to RGB for PIL
            if len(frame.shape) == 3:
                frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            else:
                frame_rgb = frame
            
            # Apply ROI cropping if specified
            if roi and self._validate_roi(roi):
                x, y, w, h = roi['x'], roi['y'], roi['width'], roi['height']
                frame_rgb = frame_rgb[y:y+h, x:x+w]
            
            # Convert to PIL Image
            pil_image = Image.fromarray(frame_rgb)
            
            # Resize for better OCR (upscale small text)
            original_size = pil_image.size
            scale_factor = max(2.0, 1200 / max(original_size))
            new_size = (int(original_size[0] * scale_factor), int(original_size[1] * scale_factor))
            pil_image = pil_image.resize(new_size, Image.Resampling.LANCZOS)
            
            # Convert to grayscale
            pil_image = pil_image.convert('L')
            
            # Enhance contrast for text visibility
            enhancer = ImageEnhance.Contrast(pil_image)
            pil_image = enhancer.enhance(2.0)
            
            # Apply sharpening filter
            pil_image = pil_image.filter(ImageFilter.SHARPEN)
            
            # Threshold to clean binary image
            img_array = np.array(pil_image)
            _, binary = cv2.threshold(img_array, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
            
            # Remove noise with morphological operations
            kernel = np.ones((2,2), np.uint8)
            cleaned = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel)
            cleaned = cv2.morphologyEx(cleaned, cv2.MORPH_OPEN, kernel)
            
            return Image.fromarray(cleaned)
            
        except Exception as e:
            logger.error(f"Frame preprocessing failed: {e}")
            # Fallback to simple conversion
            if len(frame.shape) == 3:
                return Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)).convert('L')
            else:
                return Image.fromarray(frame)
    
    def extract_text_from_frame(
        self, 
        frame: np.ndarray, 
        roi: Optional[Dict[str, Any]] = None,
        high_quality: bool = True
    ) -> Tuple[str, float, Dict[str, Any]]:
        """
        Extract text from a video frame using pre-loaded Tesseract model.
        
        Args:
            frame: Input frame as numpy array
            roi: Optional region of interest for OCR
            high_quality: Use high-quality OCR settings
            
        Returns:
            Tuple of (extracted_text, confidence_score, ocr_data)
        """
        if not self.api:
            logger.error("TesseOCR API not initialized")
            return "", 0.0, {}
            
        start_time = time.time()
        
        try:
            with self._lock:  # Thread safety
                # Preprocess frame
                processed_image = self.preprocess_frame_for_ocr(frame, roi)
                
                # Set DPI for better recognition
                dpi = self.frame_config['high_quality_dpi'] if high_quality else self.frame_config['dpi']
                
                # Set image in Tesseract API
                self.api.SetImage(processed_image)
                
                # Extract text - this is much faster than pytesseract!
                extracted_text = self.api.GetUTF8Text().strip()
                
                # Get confidence score
                confidence = self.api.MeanTextConf()
                confidence_normalized = confidence / 100.0 if confidence > 0 else 0.0
                
                # Get detailed word-level information if needed
                ocr_data = {}
                try:
                    # Get bounding boxes and confidences for each word
                    boxes = self.api.GetComponentImages(tesserocr.RIL.WORD, True)
                    words_info = []
                    
                    for (box, text, conf) in boxes:
                        if text.strip() and conf > (self.frame_config['high_quality_min_confidence'] 
                                                   if high_quality else self.frame_config['min_confidence']):
                            words_info.append({
                                'text': text.strip(),
                                'confidence': conf,
                                'bbox': box
                            })
                    
                    ocr_data = {
                        'words': words_info,
                        'processing_time': time.time() - start_time,
                        'dpi': dpi,
                        'method': 'tesserocr'
                    }
                    
                except Exception as e:
                    logger.debug(f"Could not extract detailed OCR data: {e}")
            
            # Update performance metrics
            self.processed_frames += 1
            self.total_processing_time += time.time() - start_time
            
            logger.debug(f"TesseOCR extracted text: '{extracted_text[:50]}...' "
                        f"with confidence {confidence_normalized:.2f} "
                        f"in {time.time() - start_time:.3f}s")
            
            return extracted_text, confidence_normalized, ocr_data
            
        except Exception as e:
            logger.error(f"TesseOCR text extraction failed: {e}")
            return "", 0.0, {}
    
    def extract_text_from_frame_batch(
        self, 
        frames: List[np.ndarray], 
        roi: Optional[Dict[str, Any]] = None,
        high_quality: bool = True
    ) -> List[Tuple[str, float, Dict[str, Any]]]:
        """
        Process multiple frames in batch for maximum efficiency.
        
        Args:
            frames: List of input frames as numpy arrays
            roi: Optional region of interest for OCR
            high_quality: Use high-quality OCR settings
            
        Returns:
            List of tuples (extracted_text, confidence_score, ocr_data)
        """
        if not self.api:
            logger.error("TesseOCR API not initialized")
            return [("", 0.0, {})] * len(frames)
        
        results = []
        batch_start = time.time()
        
        logger.info(f"Processing batch of {len(frames)} frames with TesseOCR...")
        
        for i, frame in enumerate(frames):
            try:
                result = self.extract_text_from_frame(frame, roi, high_quality)
                results.append(result)
                
                if i % 10 == 0 and i > 0:
                    avg_time = (time.time() - batch_start) / (i + 1)
                    logger.debug(f"Processed {i+1}/{len(frames)} frames, "
                               f"avg time: {avg_time:.3f}s/frame")
                    
            except Exception as e:
                logger.error(f"Error processing frame {i}: {e}")
                results.append(("", 0.0, {}))
        
        total_time = time.time() - batch_start
        logger.info(f"Batch processing complete: {len(frames)} frames in {total_time:.2f}s "
                   f"({total_time/len(frames):.3f}s/frame)")
        
        return results
    
    def detect_sensitive_content(
        self, 
        frame: np.ndarray, 
        roi: Optional[Dict[str, Any]] = None
    ) -> Tuple[bool, str, Dict[str, Any]]:
        """
        Detect sensitive medical information in frame using TesseOCR.
        
        Args:
            frame: Input frame as numpy array
            roi: Optional region of interest
            
        Returns:
            Tuple of (has_sensitive_content, extracted_text, pattern_matches)
        """
        try:
            # Extract text using high-quality OCR
            ocr_text, confidence, ocr_data = self.extract_text_from_frame(
                frame, roi, high_quality=True
            )
            
            if not ocr_text or confidence < 0.3:
                return False, "", {}
            
            # Extract medical patterns
            patterns = self.extract_medical_text_patterns(ocr_text)
            
            # Determine if sensitive content is present
            has_sensitive = bool(patterns)
            
            logger.debug(f"Sensitive content detection: {has_sensitive}, "
                        f"patterns found: {len(patterns)}")
            
            return has_sensitive, ocr_text, patterns
            
        except Exception as e:
            logger.error(f"Sensitive content detection failed: {e}")
            return False, "", {}
    
    def extract_medical_text_patterns(self, text: str) -> Dict[str, Any]:
        """
        Extract medical information patterns from text.
        
        Args:
            text: Input text to analyze
            
        Returns:
            Dictionary of extracted medical information
        """
        extracted_info = {
            'patient_name': [],
            'birth_date': [],
            'case_number': [],
            'examination_date': [],
            'gender': [],
            'dates': [],
            'numbers': []
        }
        
        try:
            for pattern in self.medical_patterns:
                matches = re.finditer(pattern, text, re.IGNORECASE)
                for match in matches:
                    matched_text = match.group(1) if match.groups() else match.group(0)
                    
                    # Categorize matches
                    if any(keyword in pattern.lower() for keyword in ['patient', 'pat']):
                        extracted_info['patient_name'].append(matched_text.strip())
                    elif 'geb' in pattern.lower() or 'geburtsdatum' in pattern.lower():
                        extracted_info['birth_date'].append(matched_text.strip())
                    elif 'fall' in pattern.lower():
                        extracted_info['case_number'].append(matched_text.strip())
                    elif 'datum' in pattern.lower():
                        extracted_info['examination_date'].append(matched_text.strip())
                    elif 'geschlecht' in pattern.lower():
                        extracted_info['gender'].append(matched_text.strip())
                    elif re.match(r'\d{1,2}\.\d{1,2}\.\d{2,4}', matched_text):
                        extracted_info['dates'].append(matched_text.strip())
                    elif re.match(r'\d{2,}', matched_text):
                        extracted_info['numbers'].append(matched_text.strip())
            
            # Remove duplicates
            for key in extracted_info:
                extracted_info[key] = list(set(extracted_info[key]))
            
            return extracted_info
            
        except Exception as e:
            logger.error(f"Medical pattern extraction failed: {e}")
            return extracted_info
    
    def _validate_roi(self, roi: Dict[str, Any]) -> bool:
        """Validate ROI parameters."""
        required_keys = ['x', 'y', 'width', 'height']
        return all(key in roi for key in required_keys) and all(
            isinstance(roi[key], (int, float)) and roi[key] >= 0 
            for key in required_keys
        )
    
    def get_performance_stats(self) -> Dict[str, float]:
        """Get performance statistics."""
        if self.processed_frames == 0:
            return {'avg_processing_time': 0.0, 'total_frames': 0}
        
        return {
            'avg_processing_time': self.total_processing_time / self.processed_frames,
            'total_frames': self.processed_frames,
            'total_time': self.total_processing_time,
            'frames_per_second': self.processed_frames / self.total_processing_time if self.total_processing_time > 0 else 0
        }
    
    def reset_performance_stats(self):
        """Reset performance metrics."""
        self.processed_frames = 0
        self.total_processing_time = 0.0


# Global instance for easy access (optional - can be instantiated as needed)
_global_tesseocr_processor = None

def get_global_tesseocr_processor(language: str = 'deu+eng') -> TesseOCRFrameProcessor:
    """
    Get or create global TesseOCR processor instance.
    
    This allows for model reuse across multiple video processing sessions.
    """
    global _global_tesseocr_processor
    
    if _global_tesseocr_processor is None:
        _global_tesseocr_processor = TesseOCRFrameProcessor(language)
    
    return _global_tesseocr_processor


# Convenience function for drop-in replacement
def extract_text_from_frame_fast(
    frame: np.ndarray, 
    roi: Optional[Dict[str, Any]] = None,
    high_quality: bool = True,
    language: str = 'deu+eng'
) -> Tuple[str, float, Dict[str, Any]]:
    """
    Fast text extraction from frame using global TesseOCR processor.
    
    Drop-in replacement for pytesseract-based functions with significant speed improvement.
    """
    processor = get_global_tesseocr_processor(language)
    return processor.extract_text_from_frame(frame, roi, high_quality)
