import cv2
import pytesseract
from PIL import Image, ImageEnhance, ImageFilter
import numpy as np
from typing import Dict, List, Tuple
import os

class ImageAnalyzer:
    def __init__(self):
        print("🔄 Initializing image analyzer...")
        self.supported_formats = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp', '.heic', '.heif']
        print(f"   Supported formats: {self.supported_formats}")

        # Check if pillow-heif is available for HEIC support
        try:
            import pillow_heif
            pillow_heif.register_heif_opener()
            print("   ✅ HEIC support enabled")
        except ImportError:
            print("   ⚠️  HEIC support not available (install pillow-heif for HEIC files)")

        print("   ✅ Image analyzer ready")

    def analyze_image(self, image_path: str) -> Dict:
        print(f"📸 Analyzing image: {os.path.basename(image_path)}")

        if not self._is_supported_format(image_path):
            print(f"   ❌ Unsupported format")
            return {"error": "Unsupported image format"}

        try:
            print("   Loading image...")
            # Load image
            image = Image.open(image_path)
            cv_image = cv2.imread(image_path)
            print(f"   ✅ Image loaded - Size: {image.size}, Format: {image.format}")

            print("   Extracting text...")
            text_info = self._extract_text(image)
            print(f"   📝 Text extraction complete - Found {text_info['word_count']} words")

            print("   Analyzing visual elements...")
            visual_elements = self._analyze_visual_elements(cv_image)

            print("   Analyzing colors...")
            color_analysis = self._analyze_colors(cv_image)

            print("   Checking quality metrics...")
            quality_metrics = self._analyze_quality(cv_image)

            analysis = {
                "file_path": image_path,
                "file_size_mb": os.path.getsize(image_path) / (1024 * 1024),
                "dimensions": image.size,
                "format": image.format,
                "mode": image.mode,
                "extracted_text": text_info,
                "visual_elements": visual_elements,
                "color_analysis": color_analysis,
                "quality_metrics": quality_metrics
            }

            print("   ✅ Image analysis complete")
            return analysis

        except Exception as e:
            print(f"   ❌ Analysis failed: {str(e)}")
            return {"error": f"Failed to analyze image: {str(e)}"}

    def _is_supported_format(self, image_path: str) -> bool:
        _, ext = os.path.splitext(image_path.lower())
        return ext in self.supported_formats

    def _extract_text(self, image: Image) -> Dict:
        try:
            print("     Enhancing image for OCR...")
            # Enhance image for better OCR
            enhanced = ImageEnhance.Contrast(image).enhance(2.0)
            enhanced = ImageEnhance.Sharpness(enhanced).enhance(2.0)

            print("     Running OCR...")
            # Extract text
            text = pytesseract.image_to_string(enhanced).strip()

            # Get confidence scores
            data = pytesseract.image_to_data(enhanced, output_type=pytesseract.Output.DICT)
            confidences = [int(conf) for conf in data['conf'] if int(conf) > 0]
            avg_confidence = sum(confidences) / len(confidences) if confidences else 0

            print(f"     OCR confidence: {avg_confidence:.1f}%")

            return {
                "text": text,
                "word_count": len(text.split()) if text else 0,
                "confidence": avg_confidence,
                "has_text": bool(text and len(text.strip()) > 0)
            }
        except Exception as e:
            print(f"     ⚠️  OCR failed: {e}")
            return {
                "text": "",
                "word_count": 0,
                "confidence": 0,
                "has_text": False
            }

    def _analyze_visual_elements(self, cv_image: np.ndarray) -> Dict:
        gray = cv2.cvtColor(cv_image, cv2.COLOR_BGR2GRAY)

        # Detect edges
        edges = cv2.Canny(gray, 50, 150)
        edge_density = np.sum(edges > 0) / edges.size

        # Detect shapes/contours
        contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        # Analyze contour complexity
        complex_shapes = 0
        for contour in contours:
            if cv2.contourArea(contour) > 100:
                complex_shapes += 1

        return {
            "edge_density": float(edge_density),
            "num_shapes": len(contours),
            "complex_shapes": complex_shapes,
            "has_geometric_elements": complex_shapes > 3
        }

    def _analyze_colors(self, cv_image: np.ndarray) -> Dict:
        # Convert to RGB
        rgb_image = cv2.cvtColor(cv_image, cv2.COLOR_BGR2RGB)

        # Calculate dominant colors
        pixels = rgb_image.reshape(-1, 3)

        # Simple color analysis
        mean_color = np.mean(pixels, axis=0)
        color_variance = np.var(pixels, axis=0)

        # Check for vibrant colors
        vibrant_pixels = np.sum(np.max(pixels, axis=1) > 200)
        vibrant_ratio = vibrant_pixels / len(pixels)

        return {
            "mean_color": mean_color.tolist(),
            "color_variance": color_variance.tolist(),
            "vibrant_ratio": float(vibrant_ratio),
            "is_colorful": vibrant_ratio > 0.1
        }

    def _analyze_quality(self, cv_image: np.ndarray) -> Dict:
        gray = cv2.cvtColor(cv_image, cv2.COLOR_BGR2GRAY)

        # Brightness
        brightness = np.mean(gray)

        # Contrast (standard deviation)
        contrast = np.std(gray)

        # Sharpness (Laplacian variance)
        sharpness = cv2.Laplacian(gray, cv2.CV_64F).var()

        # Noise level (estimate using edge detection)
        edges = cv2.Canny(gray, 50, 150)
        noise_level = np.sum(edges > 0) / edges.size

        # More lenient quality thresholds for real-world content
        is_high_quality = sharpness > 50 and contrast > 15  # Much more lenient
        print(f"     Quality: Sharpness={sharpness:.1f}, Contrast={contrast:.1f} -> {'High' if is_high_quality else 'Low'}")

        return {
            "brightness": float(brightness),
            "contrast": float(contrast),
            "sharpness": float(sharpness),
            "noise_level": float(noise_level),
            "is_high_quality": is_high_quality
        }