"""
Image Extractor utility for extracting embedded images from documents.
Page snapshots are not generated; preview uses PDF rendering instead.
"""
import os
from pathlib import Path
from typing import List, Dict
from PIL import Image


class ImageExtractor:
    """Extract embedded images from PDF, DOCX, and standalone image files."""

    def __init__(self, output_folder: str = "extracted_images"):
        self.output_folder = output_folder
        Path(output_folder).mkdir(parents=True, exist_ok=True)
        self._image_counter = {}

    def _get_unique_image_id(self, file_id: str, base_name: str) -> str:
        key = f"{file_id}_{base_name}"
        if key not in self._image_counter:
            self._image_counter[key] = 0
        self._image_counter[key] += 1
        return f"{base_name}_{self._image_counter[key]}"

    def extract_from_pdf(self, file_path: str, file_id: str,
                         high_quality: bool = True) -> List[Dict]:
        """Extract embedded images from PDF (no page snapshots)."""
        extracted_images = []

        try:
            import pymupdf
            doc = pymupdf.open(file_path)

            for page_num in range(len(doc)):
                page = doc[page_num]
                image_list = page.get_images()

                for img_index, img in enumerate(image_list):
                    try:
                        xref = img[0]
                        base_image = doc.extract_image(xref)
                        image_bytes = base_image["image"]
                        image_ext = base_image["ext"]

                        unique_id = self._get_unique_image_id(
                            file_id, f"p{page_num + 1}_img{img_index + 1}"
                        )
                        img_filename = f"{file_id}_{unique_id}.{image_ext}"
                        img_path = os.path.join(self.output_folder, img_filename)

                        with open(img_path, "wb") as img_file:
                            img_file.write(image_bytes)

                        extracted_images.append({
                            'path': img_path,
                            'page': page_num + 1,
                            'type': 'embedded_image',
                            'index': img_index + 1
                        })
                    except Exception as e:
                        print(f"Error extracting embedded image on page {page_num + 1}: {e}")
                        continue

            doc.close()

        except ImportError:
            print("pymupdf is required for PDF image extraction")
        except Exception as e:
            print(f"Error opening PDF for image extraction: {e}")

        return extracted_images

    def extract_from_docx(self, file_path: str, file_id: str) -> List[Dict]:
        """Extract embedded images from DOCX (no page snapshots)."""
        extracted_images = []

        try:
            import docx
            doc = docx.Document(file_path)
            img_index = 0

            for rel in doc.part.rels.values():
                if "image" in rel.target_ref:
                    try:
                        img_data = rel.target_part.blob
                        ext = rel.target_ref.split('.')[-1]
                        img_index += 1

                        unique_id = self._get_unique_image_id(file_id, f"docx_img{img_index}")
                        img_filename = f"{file_id}_{unique_id}.{ext}"
                        img_path = os.path.join(self.output_folder, img_filename)

                        with open(img_path, "wb") as img_file:
                            img_file.write(img_data)

                        img_verify = Image.open(img_path)
                        img_verify.verify()

                        extracted_images.append({
                            'path': img_path,
                            'page': 1,
                            'type': 'docx_image',
                            'index': img_index,
                            'size': len(img_data)
                        })
                    except Exception as e:
                        print(f"Error extracting DOCX image {img_index}: {e}")
                        continue

        except ImportError:
            print("python-docx is required for DOCX image extraction")
        except Exception as e:
            print(f"Error opening DOCX for image extraction: {e}")

        return extracted_images

    def extract_from_image_file(self, file_path: str, file_id: str) -> List[Dict]:
        """Process standalone image files."""
        extracted_images = []

        try:
            img = Image.open(file_path)
            img.verify()
            img = Image.open(file_path)

            ext = file_path.split('.')[-1].lower()
            unique_id = self._get_unique_image_id(file_id, "image")
            img_filename = f"{file_id}_{unique_id}.{ext}"
            img_path = os.path.join(self.output_folder, img_filename)

            if img.mode not in ('RGB', 'L'):
                img = img.convert('RGB')
            img.save(img_path, format='PNG', optimize=True)

            extracted_images.append({
                'path': img_path,
                'page': 1,
                'type': 'standalone_image',
                'index': 1,
                'original_path': file_path
            })

        except Exception as e:
            print(f"Error processing image file: {e}")

        return extracted_images
