"""
DocRes Image Restoration Engine

This module provides a wrapper around the DocRes inference functionality
for easy integration with Doctra's document processing pipeline.

DocRes supports 5 restoration tasks:
- dewarping: Corrects document perspective distortion
- deshadowing: Removes shadows from documents  
- appearance: General appearance enhancement
- deblurring: Reduces blur in document images
- binarization: Converts to clean black/white text
- end2end: Pipeline combining dewarping → deshadowing → appearance
"""

import os
import sys
import cv2
import numpy as np
import torch
import tempfile
import time
from pathlib import Path
from typing import Union, List, Tuple, Optional, Dict, Any

# Hugging Face Hub imports
try:
    from huggingface_hub import hf_hub_download
    from huggingface_hub.utils import disable_progress_bars
    disable_progress_bars()
    HF_HUB_AVAILABLE = True
except ImportError:
    HF_HUB_AVAILABLE = False

# Progress bar imports
from doctra.utils.progress import create_beautiful_progress_bar, create_notebook_friendly_bar

# Add DocRes to path and change to DocRes directory for relative imports
current_dir = Path(__file__).parent
docres_dir = current_dir.parent.parent / "third_party" / "docres"
sys.path.insert(0, str(docres_dir))

# Store original working directory
original_cwd = os.getcwd()

try:
    # Change to DocRes directory for relative imports to work
    os.chdir(str(docres_dir))
    
    # Now import DocRes modules (they use relative imports)
    from inference import (
        model_init, inference_one_im, dewarping, deshadowing, 
        appearance, deblurring, binarization
    )
    from utils import convert_state_dict
    from models import restormer_arch
    from data.preprocess.crop_merge_image import stride_integral
    from data.MBD.infer import net1_net2_infer_single_im
    
    DOCRES_AVAILABLE = True
except ImportError as e:
    DOCRES_AVAILABLE = False
    # Don't print warning here, let the user handle it when they try to use it
finally:
    # Always restore original working directory
    os.chdir(original_cwd)


def load_docres_weights_from_hf():
    """
    Load DocRes model weights from Hugging Face Hub.
    
    Returns:
        Tuple of (mbd_path, docres_path) - paths to downloaded model files
    """
    if not HF_HUB_AVAILABLE:
        raise ImportError(
            "huggingface_hub is required for downloading models from Hugging Face. "
            "Install with: pip install huggingface_hub"
        )
    
    try:
        # Detect environment for progress bar
        is_notebook = "ipykernel" in sys.modules or "jupyter" in sys.modules
        
        # Create progress bar for model downloads
        if is_notebook:
            progress_bar = create_notebook_friendly_bar(
                total=2, 
                desc="Downloading DocRes models from Hugging Face Hub"
            )
        else:
            progress_bar = create_beautiful_progress_bar(
                total=2, 
                desc="Downloading DocRes models from Hugging Face Hub",
                leave=True
            )
        
        with progress_bar:
            # Download DocRes main model
            _ = hf_hub_download("DaVinciCode/doctra-docres-main", filename="config.json")
            docres_path = hf_hub_download("DaVinciCode/doctra-docres-main", filename="docres.pkl")
            progress_bar.update(1)
            
            # Download MBD model
            _ = hf_hub_download("DaVinciCode/doctra-docres-mbd", filename="config.json")
            mbd_path = hf_hub_download("DaVinciCode/doctra-docres-mbd", filename="mbd.pkl")
            progress_bar.update(1)
        
        # Verify file sizes (silently)
        docres_size = Path(docres_path).stat().st_size
        mbd_size = Path(mbd_path).stat().st_size
        
        return mbd_path, docres_path
        
    except Exception as e:
        raise RuntimeError(f"Failed to download models from Hugging Face: {e}")


def get_model_paths(use_huggingface: bool = True, model_path: Optional[str] = None, mbd_path: Optional[str] = None):
    """
    Get model paths, either from Hugging Face or local files.
    
    Args:
        use_huggingface: Whether to use Hugging Face Hub for model loading
        model_path: Local path to DocRes model (if not using Hugging Face)
        mbd_path: Local path to MBD model (if not using Hugging Face)
        
    Returns:
        Tuple of (mbd_path, docres_path)
    """
    if use_huggingface and HF_HUB_AVAILABLE:
        try:
            return load_docres_weights_from_hf()
        except Exception as e:
            print(f"⚠️ Hugging Face download failed: {e}")
            print("   Falling back to local model files...")
            use_huggingface = False
    
    if not use_huggingface:
        # Use local model files
        if model_path is None:
            model_path = docres_dir / "checkpoints" / "docres.pkl"
        if mbd_path is None:
            mbd_path = docres_dir / "data" / "MBD" / "checkpoint" / "mbd.pkl"
        
        return str(mbd_path), str(model_path)
    
    raise RuntimeError("Cannot load models: Hugging Face Hub not available and no local paths provided")


class DocResEngine:
    """
    DocRes Image Restoration Engine
    
    A wrapper around DocRes inference functionality for easy integration
    with Doctra's document processing pipeline.
    """
    
    SUPPORTED_TASKS = [
        'dewarping', 'deshadowing', 'appearance', 
        'deblurring', 'binarization', 'end2end'
    ]
    
    def __init__(
        self, 
        device: Optional[str] = None,
        use_half_precision: bool = True,
        model_path: Optional[str] = None,
        mbd_path: Optional[str] = None
    ):
        """
        Initialize DocRes Engine
        
        Args:
            device: Device to run on ('cuda', 'cpu', or None for auto-detect)
            use_half_precision: Whether to use half precision for inference
            model_path: Path to DocRes model checkpoint (optional, defaults to Hugging Face Hub)
            mbd_path: Path to MBD model checkpoint (optional, defaults to Hugging Face Hub)
        """
        if not DOCRES_AVAILABLE:
            raise ImportError(
                "DocRes is not available. Please install the missing dependencies:\n"
                "pip install scikit-image>=0.19.3\n\n"
                "The DocRes module is already included in this library, but requires "
                "scikit-image for image processing operations."
            )
        
        # Set device
        if device is None:
            self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        else:
            requested_device = torch.device(device)
            # Check if the requested device is available
            if requested_device.type == 'cuda' and not torch.cuda.is_available():
                print(f"Warning: CUDA requested but not available. Falling back to CPU.")
                self.device = torch.device('cpu')
            else:
                self.device = requested_device
        
        self.use_half_precision = use_half_precision
        
        # Get model paths (always from Hugging Face Hub)
        try:
            self.mbd_path, self.model_path = get_model_paths(
                use_huggingface=True,
                model_path=model_path,
                mbd_path=mbd_path
            )
        except Exception as e:
            raise RuntimeError(f"Failed to get model paths: {e}")
        
        # Verify model files exist
        if not os.path.exists(self.model_path):
            raise FileNotFoundError(
                f"DocRes model not found at {self.model_path}. "
                f"This may indicate a Hugging Face download failure. "
                f"Please check your internet connection and try again."
            )
        
        if not os.path.exists(self.mbd_path):
            raise FileNotFoundError(
                f"MBD model not found at {self.mbd_path}. "
                f"This may indicate a Hugging Face download failure. "
                f"Please check your internet connection and try again."
            )
        
        # Initialize model
        self._model = None
        self._initialize_model()
    
    def _initialize_model(self):
        """Initialize the DocRes model"""
        try:
            # Create model architecture
            self._model = restormer_arch.Restormer( 
                inp_channels=6, 
                out_channels=3, 
                dim=48,
                num_blocks=[2,3,3,4], 
                num_refinement_blocks=4,
                heads=[1,2,4,8],
                ffn_expansion_factor=2.66,
                bias=False,
                LayerNorm_type='WithBias',
                dual_pixel_task=True        
            )
            
            # Load model weights - always load to CPU first, then move to target device
            state = convert_state_dict(torch.load(self.model_path, map_location='cpu')['model_state'])
            
            self._model.load_state_dict(state)
            self._model.eval()
            self._model = self._model.to(self.device)
            
        except Exception as e:
            raise RuntimeError(f"Failed to initialize DocRes model: {e}")
    
    def restore_image(
        self, 
        image: Union[str, np.ndarray], 
        task: str = "appearance",
        save_prompts: bool = False
    ) -> Tuple[np.ndarray, Dict[str, Any]]:
        """
        Restore a single image using DocRes
        
        Args:
            image: Path to image file or numpy array
            task: Restoration task to perform
            save_prompts: Whether to save intermediate prompts
            
        Returns:
            Tuple of (restored_image, metadata)
        """
        if task not in self.SUPPORTED_TASKS:
            raise ValueError(f"Unsupported task: {task}. Supported tasks: {self.SUPPORTED_TASKS}")
        
        # Load image if path provided
        if isinstance(image, str):
            if not os.path.exists(image):
                raise FileNotFoundError(f"Image not found: {image}")
            img_array = cv2.imread(image)
            if img_array is None:
                raise ValueError(f"Could not load image: {image}")
        else:
            img_array = image.copy()
        
        original_shape = img_array.shape
        
        try:
            # Handle end2end pipeline
            if task == "end2end":
                return self._run_end2end_pipeline(img_array, save_prompts)
            
            # Run single task
            restored_img, metadata = self._run_single_task(img_array, task, save_prompts)
            
            metadata.update({
                'original_shape': original_shape,
                'restored_shape': restored_img.shape,
                'task': task,
                'device': str(self.device)
            })
            
            return restored_img, metadata
            
        except Exception as e:
            raise RuntimeError(f"Image restoration failed: {e}")
    
    def _run_single_task(self, img_array: np.ndarray, task: str, save_prompts: bool) -> Tuple[np.ndarray, Dict]:
        """Run a single restoration task"""
        
        # Create temporary file for inference
        with tempfile.NamedTemporaryFile(suffix='.jpg', delete=False) as tmp_file:
            tmp_path = tmp_file.name
            cv2.imwrite(tmp_path, img_array)
        
        try:
            # Change to DocRes directory for inference to work properly
            original_cwd = os.getcwd()
            os.chdir(str(docres_dir))
            
            # Set global DEVICE variable that DocRes inference expects
            import inference  # Import the inference module to set its global DEVICE
            inference.DEVICE = self.device
            
            try:
                # Run inference
                prompt1, prompt2, prompt3, restored = inference_one_im(self._model, tmp_path, task)
            finally:
                # Always restore original working directory
                os.chdir(original_cwd)
            
            metadata = {
                'task': task,
                'device': str(self.device)
            }
            
            if save_prompts:
                metadata['prompts'] = {
                    'prompt1': prompt1,
                    'prompt2': prompt2, 
                    'prompt3': prompt3
                }
            
            return restored, metadata
            
        finally:
            # Clean up temporary file with retry for Windows
            try:
                # Wait a bit for file handles to be released
                time.sleep(0.1)
                os.unlink(tmp_path)
            except PermissionError:
                # If still locked, try again after a longer wait
                time.sleep(1)
                try:
                    os.unlink(tmp_path)
                except PermissionError:
                    # If still failing, just leave it - it will be cleaned up by the OS
                    pass
    
    def _run_end2end_pipeline(self, img_array: np.ndarray, save_prompts: bool) -> Tuple[np.ndarray, Dict]:
        """Run the end2end pipeline: dewarping → deshadowing → appearance"""
        
        intermediate_steps = {}
        
        # Change to DocRes directory for inference to work properly
        original_cwd = os.getcwd()
        os.chdir(str(docres_dir))
        
        # Set global DEVICE variable that DocRes inference expects
        import inference  # Import the inference module to set its global DEVICE
        inference.DEVICE = self.device
        
        try:
            with tempfile.TemporaryDirectory() as tmp_dir:
                # Step 1: Dewarping
                step1_path = os.path.join(tmp_dir, "step1.jpg")
                cv2.imwrite(step1_path, img_array)
                
                prompt1, prompt2, prompt3, dewarped = inference_one_im(self._model, step1_path, "dewarping")
                intermediate_steps['dewarped'] = dewarped
                
                # Step 2: Deshadowing
                step2_path = os.path.join(tmp_dir, "step2.jpg")
                cv2.imwrite(step2_path, dewarped)
                
                prompt1, prompt2, prompt3, deshadowed = inference_one_im(self._model, step2_path, "deshadowing")
                intermediate_steps['deshadowed'] = deshadowed
                
                # Step 3: Appearance
                step3_path = os.path.join(tmp_dir, "step3.jpg")
                cv2.imwrite(step3_path, deshadowed)
                
                prompt1, prompt2, prompt3, final = inference_one_im(self._model, step3_path, "appearance")
                
                metadata = {
                    'task': 'end2end',
                    'device': str(self.device),
                    'intermediate_steps': intermediate_steps
                }
                
                if save_prompts:
                    metadata['prompts'] = {
                        'prompt1': prompt1,
                        'prompt2': prompt2,
                        'prompt3': prompt3
                    }
                
                return final, metadata
        finally:
            # Always restore original working directory
            os.chdir(original_cwd)
    
    def batch_restore(
        self, 
        images: List[Union[str, np.ndarray]], 
        task: str = "appearance",
        save_prompts: bool = False
    ) -> List[Tuple[Optional[np.ndarray], Dict[str, Any]]]:
        """
        Restore multiple images in batch
        
        Args:
            images: List of image paths or numpy arrays
            task: Restoration task to perform
            save_prompts: Whether to save intermediate prompts
            
        Returns:
            List of (restored_image, metadata) tuples
        """
        results = []
        
        for i, image in enumerate(images):
            try:
                restored_img, metadata = self.restore_image(image, task, save_prompts)
                results.append((restored_img, metadata))
            except Exception as e:
                # Return None for failed images with error metadata
                error_metadata = {
                    'error': str(e),
                    'task': task,
                    'device': str(self.device),
                    'image_index': i
                }
                results.append((None, error_metadata))
        
        return results
    
    def get_supported_tasks(self) -> List[str]:
        """Get list of supported restoration tasks"""
        return self.SUPPORTED_TASKS.copy()
    
    def is_available(self) -> bool:
        """Check if DocRes is available and properly configured"""
        return DOCRES_AVAILABLE and self._model is not None
    
    def restore_pdf(
        self, 
        pdf_path: str, 
        output_path: str | None = None,
        task: str = "appearance",
        dpi: int = 200
    ) -> str | None:
        """
        Restore an entire PDF document using DocRes
        
        Args:
            pdf_path: Path to the input PDF file
            output_path: Path for the enhanced PDF (if None, auto-generates)
            task: DocRes restoration task (default: "appearance")
            dpi: DPI for PDF rendering (default: 200)
            
        Returns:
            Path to the enhanced PDF or None if failed
        """
        try:
            from PIL import Image
            from doctra.utils.pdf_io import render_pdf_to_images
            
            # Generate output path if not provided
            if output_path is None:
                pdf_dir = os.path.dirname(pdf_path)
                pdf_name = os.path.splitext(os.path.basename(pdf_path))[0]
                output_path = os.path.join(pdf_dir, f"{pdf_name}_enhanced.pdf")
            
            print(f"🔄 Processing PDF with DocRes: {os.path.basename(pdf_path)}")
            
            # Render all pages to images
            pil_pages = [im for (im, _, _) in render_pdf_to_images(pdf_path, dpi=dpi)]
            
            if not pil_pages:
                print("❌ No pages found in PDF")
                return None
            
            # Process each page with DocRes
            enhanced_pages = []
            
            # Detect environment for progress bar
            is_notebook = "ipykernel" in sys.modules or "jupyter" in sys.modules
            
            # Create progress bar for page processing
            if is_notebook:
                progress_bar = create_notebook_friendly_bar(
                    total=len(pil_pages), 
                    desc="Processing pages"
                )
            else:
                progress_bar = create_beautiful_progress_bar(
                    total=len(pil_pages), 
                    desc="Processing pages",
                    leave=True
                )
            
            with progress_bar:
                for i, page_img in enumerate(pil_pages):
                    try:
                        # Convert PIL to numpy array
                        img_array = np.array(page_img)
                        
                        # Apply DocRes restoration
                        restored_img, _ = self.restore_image(img_array, task)
                        
                        # Convert back to PIL Image
                        enhanced_page = Image.fromarray(restored_img)
                        enhanced_pages.append(enhanced_page)
                        
                        progress_bar.set_description(f"✅ Page {i+1}/{len(pil_pages)} processed")
                        progress_bar.update(1)
                        
                    except Exception as e:
                        print(f"  ⚠️ Page {i+1} processing failed: {e}, using original")
                        enhanced_pages.append(page_img)
                        progress_bar.set_description(f"⚠️ Page {i+1} failed, using original")
                        progress_bar.update(1)
            
            # Create enhanced PDF
            if enhanced_pages:
                enhanced_pages[0].save(
                    output_path,
                    "PDF",
                    resolution=100.0,
                    save_all=True,
                    append_images=enhanced_pages[1:] if len(enhanced_pages) > 1 else []
                )
                
                print(f"✅ Enhanced PDF saved: {output_path}")
                return output_path
            else:
                print("❌ No pages to save")
                return None
                
        except ImportError as e:
            print(f"❌ Required dependencies not available: {e}")
            print("Install with: pip install PyMuPDF")
            return None
        except Exception as e:
            print(f"❌ Error processing PDF with DocRes: {e}")
            return None