# utils.py
import os
from transformers import AutoModel, AutoTokenizer
from PIL import Image, ImageEnhance, ImageFilter
import torch
import logging
from transformers import BertTokenizer

logger = logging.getLogger(__name__)

class OCRModel:
    _instance = None
    
    def __new__(cls):
        if cls._instance is None:
            cls._instance = super(OCRModel, cls).__new__(cls)
            cls._instance.initialize()
        return cls._instance
    
    def initialize(self):
        try:
            logger.info("Initializing OCR model...")
            
            # Try different tokenizer approaches
            try:
                # First try with the standard approach
                self.tokenizer = AutoTokenizer.from_pretrained(
                    'stepfun-ai/GOT-OCR2_0',
                    trust_remote_code=True,
                    use_fast=False
                )
            except Exception as e:
                logger.warning(f"Standard tokenizer failed, trying BertTokenizer: {str(e)}")
                # Fall back to BertTokenizer if AutoTokenizer fails
                self.tokenizer = BertTokenizer.from_pretrained(
                    'stepfun-ai/GOT-OCR2_0',
                    trust_remote_code=True
                )
                
            self.model = AutoModel.from_pretrained(
                'stepfun-ai/GOT-OCR2_0',
                trust_remote_code=True,
                low_cpu_mem_usage=True,
                device_map='auto',
                use_safetensors=True
            )
            
            self.device = "cuda" if torch.cuda.is_available() else "cpu"
            self.model = self.model.eval().to(self.device)
            
            logger.info("Model initialization completed successfully")
            
        except Exception as e:
            logger.error(f"Error initializing model: {str(e)}", exc_info=True)
            raise
            
    def preprocess_image(self, image):
        """Image preprocessing to improve text recognition quality"""
        try:
            # Convert image to RGB if it is not already
            if image.mode != 'RGB':
                image = image.convert('RGB')

            # Improve contrast
            enhancer = ImageEnhance.Contrast(image)
            image = enhancer.enhance(1.5)

            # Improve Sharpness
            enhancer = ImageEnhance.Sharpness(image)
            image = enhancer.enhance(1.5)

            # Improve Brightness
            enhancer = ImageEnhance.Brightness(image)
            image = enhancer.enhance(1.2)

            # Apply a filter to soften the image a little.
            image = image.filter(ImageFilter.SMOOTH)

            return image
        except Exception as e:
            logger.error(f"Error in image preprocessing: {str(e)}", exc_info=True)
            raise

    def process_image(self, image_stream):
        try:
            logger.info("Starting image processing")
            
            # Save image temporarily because the model requires a file path.
            temp_image_path = "temp_image.jpg"
            
            # Reset the start pointer for BytesIO
            image_stream.seek(0)
            
            # Open and save the image temporarily.
            image = Image.open(image_stream).convert('RGB')
            processed_image = self.preprocess_image(image)
            processed_image.save(temp_image_path)
            
            # ocr 
            try:
                result = self.model.chat(self.tokenizer, temp_image_path, ocr_type='format')
                logger.info(f"Successfully extracted text: {result[:100]}...")
                
                # Delete temporary file
                if os.path.exists(temp_image_path):
                    os.remove(temp_image_path)
                
                return result.strip()
                
            except Exception as e:
                logger.error(f"Error in OCR processing: {str(e)}", exc_info=True)
                if os.path.exists(temp_image_path):
                    os.remove(temp_image_path)
                raise
            
        except Exception as e:
            logger.error(f"Error in image processing: {str(e)}", exc_info=True)
            return f"Error processing image: {str(e)}"