"""
Enhanced PDF Parser with Image Restoration

This module provides an enhanced PDF parser that combines the structured parsing
capabilities with DocRes image restoration for improved document processing.
"""

from __future__ import annotations
import os
import sys
import numpy as np
from typing import List, Dict, Any, Optional, Union
from doctra.engines.ocr import PytesseractOCREngine, PaddleOCREngine
from contextlib import ExitStack
from PIL import Image
from tqdm import tqdm

from doctra.parsers.structured_pdf_parser import StructuredPDFParser
from doctra.engines.image_restoration import DocResEngine
from doctra.engines.vlm.service import VLMStructuredExtractor
from doctra.utils.pdf_io import render_pdf_to_images
from doctra.utils.constants import IMAGE_SUBDIRS, EXCLUDE_LABELS
from doctra.utils.file_ops import ensure_output_dirs
from doctra.utils.progress import create_beautiful_progress_bar, create_notebook_friendly_bar
from doctra.parsers.layout_order import reading_order_key
from doctra.utils.ocr_utils import ocr_box_text
from doctra.exporters.image_saver import save_box_image
from doctra.exporters.markdown_writer import write_markdown
from doctra.exporters.html_writer import write_html, write_structured_html, render_html_table, write_html_from_lines
from doctra.exporters.excel_writer import write_structured_excel
from doctra.utils.structured_utils import to_structured_dict
from doctra.exporters.markdown_table import render_markdown_table
from doctra.parsers.split_table_detector import SplitTableMatch


class EnhancedPDFParser(StructuredPDFParser):
    """
    Enhanced PDF Parser with Image Restoration capabilities.
    
    Extends the StructuredPDFParser with DocRes image restoration to improve
    document quality before processing. This is particularly useful for:
    - Scanned documents with shadows or distortion
    - Low-quality PDFs that need enhancement
    - Documents with perspective issues
    
    :param use_image_restoration: Whether to apply DocRes image restoration (default: True)
    :param restoration_task: DocRes task to use ("dewarping", "deshadowing", "appearance", "deblurring", "binarization", "end2end", default: "appearance")
    :param restoration_device: Device for DocRes processing ("cuda", "cpu", or None for auto-detect, default: None)
    :param restoration_dpi: DPI for restoration processing (default: 200)
    :param vlm: VLM engine instance (VLMStructuredExtractor). If None, VLM processing is disabled.
    :param layout_model_name: Layout detection model name (default: "PP-DocLayout_plus-L")
    :param dpi: DPI for PDF rendering (default: 200)
    :param min_score: Minimum confidence score for layout detection (default: 0.0)
    :param ocr_engine: OCR engine instance (PytesseractOCREngine or PaddleOCREngine). 
                       If None, creates a default PytesseractOCREngine with lang="eng", psm=4, oem=3.
    :param box_separator: Separator between text boxes in output (default: "\n")
    :param merge_split_tables: Whether to detect and merge split tables (default: False)
    :param bottom_threshold_ratio: Ratio for "too close to bottom" detection (default: 0.20)
    :param top_threshold_ratio: Ratio for "too close to top" detection (default: 0.15)
    :param max_gap_ratio: Maximum allowed gap between tables (default: 0.25)
    :param column_alignment_tolerance: Pixel tolerance for column alignment (default: 10.0)
    :param min_merge_confidence: Minimum confidence score for merging (default: 0.65)
    """

    def __init__(
        self,
        *,
        use_image_restoration: bool = True,
        restoration_task: str = "appearance",
        restoration_device: Optional[str] = None,
        restoration_dpi: int = 200,
        vlm: Optional[VLMStructuredExtractor] = None,
        layout_model_name: str = "PP-DocLayout_plus-L",
        dpi: int = 200,
        min_score: float = 0.0,
        ocr_engine: Optional[Union[PytesseractOCREngine, PaddleOCREngine]] = None,
        box_separator: str = "\n",
        merge_split_tables: bool = False,
        bottom_threshold_ratio: float = 0.20,
        top_threshold_ratio: float = 0.15,
        max_gap_ratio: float = 0.25,
        column_alignment_tolerance: float = 10.0,
        min_merge_confidence: float = 0.65,
    ):
        """
        Initialize the Enhanced PDF Parser with image restoration capabilities.
        """
        super().__init__(
            vlm=vlm,
            layout_model_name=layout_model_name,
            dpi=dpi,
            min_score=min_score,
            ocr_engine=ocr_engine,
            box_separator=box_separator,
            merge_split_tables=merge_split_tables,
            bottom_threshold_ratio=bottom_threshold_ratio,
            top_threshold_ratio=top_threshold_ratio,
            max_gap_ratio=max_gap_ratio,
            column_alignment_tolerance=column_alignment_tolerance,
            min_merge_confidence=min_merge_confidence,
        )
        
        self.use_image_restoration = use_image_restoration
        self.restoration_task = restoration_task
        self.restoration_device = restoration_device
        self.restoration_dpi = restoration_dpi
        
        self.docres_engine = None
        if self.use_image_restoration:
            try:
                self.docres_engine = DocResEngine(
                    device=restoration_device,
                    use_half_precision=True
                )
                print(f"✅ DocRes engine initialized with task: {restoration_task}")
            except Exception as e:
                print(f"⚠️ DocRes initialization failed: {e}")
                print("   Continuing without image restoration...")
                self.use_image_restoration = False
                self.docres_engine = None

    def parse(self, pdf_path: str, enhanced_output_dir: str = None) -> None:
        """
        Parse a PDF document with optional image restoration.
        
        :param pdf_path: Path to the input PDF file
        :param enhanced_output_dir: Directory for enhanced images (if None, uses default)
        :return: None
        """
        pdf_filename = os.path.splitext(os.path.basename(pdf_path))[0]
        
        if enhanced_output_dir is None:
            out_dir = f"outputs/{pdf_filename}/enhanced_parse"
        else:
            out_dir = enhanced_output_dir
            
        os.makedirs(out_dir, exist_ok=True)
        ensure_output_dirs(out_dir, IMAGE_SUBDIRS)
        
        if self.use_image_restoration and self.docres_engine:
            print(f"🔄 Processing PDF with image restoration: {os.path.basename(pdf_path)}")
            enhanced_pages = self._process_pages_with_restoration(pdf_path, out_dir)
            
            enhanced_pdf_path = os.path.join(out_dir, f"{pdf_filename}_enhanced.pdf")
            try:
                self._create_enhanced_pdf_from_pages(enhanced_pages, enhanced_pdf_path)
            except Exception as e:
                print(f"⚠️ Failed to create enhanced PDF: {e}")
        else:
            print(f"🔄 Processing PDF without image restoration: {os.path.basename(pdf_path)}")
            enhanced_pages = [im for (im, _, _) in render_pdf_to_images(pdf_path, dpi=self.dpi)]
        
        print("🔍 Running layout detection on enhanced pages...")
        pages = self.layout_engine.predict_pdf(
            pdf_path, batch_size=1, layout_nms=True, dpi=self.dpi, min_score=self.min_score
        )
        
        pil_pages = enhanced_pages
        
        self._process_parsing_logic(pages, pil_pages, out_dir, pdf_filename, pdf_path)

    def _process_pages_with_restoration(self, pdf_path: str, out_dir: str) -> List[Image.Image]:
        """
        Process PDF pages with DocRes image restoration.
        
        :param pdf_path: Path to the input PDF file
        :param out_dir: Output directory for enhanced images
        :return: List of enhanced PIL images
        """
        original_pages = [im for (im, _, _) in render_pdf_to_images(pdf_path, dpi=self.restoration_dpi)]
        
        if not original_pages:
            print("❌ No pages found in PDF")
            return []
        
        is_notebook = "ipykernel" in sys.modules or "jupyter" in sys.modules
        if is_notebook:
            progress_bar = create_notebook_friendly_bar(
                total=len(original_pages), 
                desc=f"DocRes {self.restoration_task}"
            )
        else:
            progress_bar = create_beautiful_progress_bar(
                total=len(original_pages), 
                desc=f"DocRes {self.restoration_task}",
                leave=True
            )
        
        enhanced_pages = []
        enhanced_dir = os.path.join(out_dir, "enhanced_pages")
        os.makedirs(enhanced_dir, exist_ok=True)
        
        try:
            with progress_bar:
                for i, page_img in enumerate(original_pages):
                    try:
                        img_array = np.array(page_img)
                        
                        restored_img, metadata = self.docres_engine.restore_image(
                            img_array, 
                            task=self.restoration_task
                        )
                        
                        enhanced_page = Image.fromarray(restored_img)
                        enhanced_pages.append(enhanced_page)
                        
                        enhanced_path = os.path.join(enhanced_dir, f"page_{i+1:03d}_enhanced.jpg")
                        enhanced_page.save(enhanced_path, "JPEG", quality=95)
                        
                        progress_bar.set_description(f"✅ Page {i+1}/{len(original_pages)} enhanced")
                        progress_bar.update(1)
                        
                    except Exception as e:
                        print(f"  ⚠️ Page {i+1} restoration failed: {e}, using original")
                        enhanced_pages.append(page_img)
                        progress_bar.set_description(f"⚠️ Page {i+1} failed, using original")
                        progress_bar.update(1)
        
        finally:
            if hasattr(progress_bar, 'close'):
                progress_bar.close()
        
        return enhanced_pages

    def _process_parsing_logic(self, pages, pil_pages, out_dir, pdf_filename, pdf_path):
        """
        Process the parsing logic with enhanced pages.
        This is extracted from the parent class to allow customization.
        """
        split_table_matches: List[SplitTableMatch] = []
        merged_table_segments = []
        
        if self.merge_split_tables and self.split_table_detector:
            try:
                split_table_matches = self.split_table_detector.detect_split_tables(pages, pil_pages)
                if split_table_matches:
                    print(f"🔗 Detected {len(split_table_matches)} split table(s) to merge")
                for match in split_table_matches:
                    merged_table_segments.append(match.segment1)
                    merged_table_segments.append(match.segment2)
            except Exception as e:
                import traceback
                traceback.print_exc()
                split_table_matches = []
        
        fig_count = sum(sum(1 for b in p.boxes if b.label == "figure") for p in pages)
        chart_count = sum(sum(1 for b in p.boxes if b.label == "chart") for p in pages)
        table_count = sum(sum(1 for b in p.boxes if b.label == "table") for p in pages)

        md_lines: List[str] = ["# Enhanced Document Content\n"]
        html_lines: List[str] = ["<h1>Enhanced Document Content</h1>"]
        structured_items: List[Dict[str, Any]] = []
        page_content: Dict[int, List[str]] = {}

        charts_desc = "Charts (VLM → table)" if self.vlm is not None else "Charts (cropped)"
        tables_desc = "Tables (VLM → table)" if self.vlm is not None else "Tables (cropped)"
        figures_desc = "Figures (cropped)"

        with ExitStack() as stack:
            is_notebook = "ipykernel" in sys.modules or "jupyter" in sys.modules
            if is_notebook:
                charts_bar = stack.enter_context(
                    create_notebook_friendly_bar(total=chart_count, desc=charts_desc)) if chart_count else None
                tables_bar = stack.enter_context(
                    create_notebook_friendly_bar(total=table_count, desc=tables_desc)) if table_count else None
                figures_bar = stack.enter_context(
                    create_notebook_friendly_bar(total=fig_count, desc=figures_desc)) if fig_count else None
            else:
                charts_bar = stack.enter_context(
                    create_beautiful_progress_bar(total=chart_count, desc=charts_desc, leave=True)) if chart_count else None
                tables_bar = stack.enter_context(
                    create_beautiful_progress_bar(total=table_count, desc=tables_desc, leave=True)) if table_count else None
                figures_bar = stack.enter_context(
                    create_beautiful_progress_bar(total=fig_count,                     desc=figures_desc, leave=True)) if fig_count else None

            for page_num in range(1, len(pil_pages) + 1):
                page_content[page_num] = [f"# Page {page_num} Content\n"]
            
            for p in pages:
                page_num = p.page_index
                page_img: Image.Image = pil_pages[page_num - 1]
                md_lines.append(f"\n## Page {page_num}\n")
                html_lines.append(f"<h2>Page {page_num}</h2>")

                for i, box in enumerate(sorted(p.boxes, key=reading_order_key), start=1):
                    if box.label in EXCLUDE_LABELS:
                        img_path = save_box_image(page_img, box, out_dir, page_num, i, IMAGE_SUBDIRS)
                        abs_img_path = os.path.abspath(img_path)
                        rel = os.path.relpath(abs_img_path, out_dir)

                        if box.label == "figure":
                            figure_md = f"![Figure — page {page_num}]({rel})\n"
                            figure_html = f'<img src="{rel}" alt="Figure — page {page_num}" />'
                            md_lines.append(figure_md)
                            html_lines.append(figure_html)
                            page_content[page_num].append(figure_md)
                            if figures_bar: figures_bar.update(1)

                        elif box.label == "chart":
                            if self.vlm is not None:
                                wrote_table = False
                                try:
                                    chart = self.vlm.extract_chart(abs_img_path)
                                    item = to_structured_dict(chart)
                                    if item:
                                        item["page"] = page_num
                                        item["type"] = "Chart"
                                        structured_items.append(item)
                                        
                                        table_md = render_markdown_table(item.get("headers"), item.get("rows"),
                                                                         title=item.get("title"))
                                        table_html = render_html_table(item.get("headers"), item.get("rows"),
                                                                       title=item.get("title"))
                                        
                                        md_lines.append(table_md)
                                        html_lines.append(table_html)
                                        page_content[page_num].append(table_md)
                                        wrote_table = True
                                except Exception as e:
                                    pass
                                if not wrote_table:
                                    chart_md = f"![Chart — page {page_num}]({rel})\n"
                                    chart_html = f'<img src="{rel}" alt="Chart — page {page_num}" />'
                                    md_lines.append(chart_md)
                                    html_lines.append(chart_html)
                                    page_content[page_num].append(chart_md)
                            else:
                                chart_md = f"![Chart — page {page_num}]({rel})\n"
                                chart_html = f'<img src="{rel}" alt="Chart — page {page_num}" />'
                                md_lines.append(chart_md)
                                html_lines.append(chart_html)
                                page_content[page_num].append(chart_md)
                            if charts_bar: charts_bar.update(1)

                        elif box.label == "table":
                            is_merged = any(seg.match_box(box, page_num) for seg in merged_table_segments)
                            if is_merged:
                                continue
                            
                            if self.vlm is not None:
                                wrote_table = False
                                try:
                                    table = self.vlm.extract_table(abs_img_path)
                                    item = to_structured_dict(table)
                                    if item:
                                        item["page"] = page_num
                                        item["type"] = "Table"
                                        structured_items.append(item)
                                        
                                        table_md = render_markdown_table(item.get("headers"), item.get("rows"),
                                                                         title=item.get("title"))
                                        table_html = render_html_table(item.get("headers"), item.get("rows"),
                                                                       title=item.get("title"))
                                        
                                        md_lines.append(table_md)
                                        html_lines.append(table_html)
                                        page_content[page_num].append(table_md)
                                        wrote_table = True
                                except Exception as e:
                                    pass
                                if not wrote_table:
                                    table_md = f"![Table — page {page_num}]({rel})\n"
                                    table_html = f'<img src="{rel}" alt="Table — page {page_num}" />'
                                    md_lines.append(table_md)
                                    html_lines.append(table_html)
                                    page_content[page_num].append(table_md)
                            else:
                                table_md = f"![Table — page {page_num}]({rel})\n"
                                table_html = f'<img src="{rel}" alt="Table — page {page_num}" />'
                                md_lines.append(table_md)
                                html_lines.append(table_html)
                                page_content[page_num].append(table_md)
                            if tables_bar: tables_bar.update(1)
                    else:
                        text = ocr_box_text(self.ocr_engine, page_img, box)
                        if text:
                            md_lines.append(text)
                            md_lines.append(self.box_separator if self.box_separator else "")
                            html_text = text.replace('\n', '<br>')
                            html_lines.append(f"<p>{html_text}</p>")
                            if self.box_separator:
                                html_lines.append("<br>")
                            page_content[page_num].append(text)
                            page_content[page_num].append(self.box_separator if self.box_separator else "")

            if split_table_matches and self.split_table_detector:
                for match_idx, match in enumerate(split_table_matches):
                    try:
                        merged_img = self.split_table_detector.merge_table_images(match)
                        
                        tables_dir = os.path.join(out_dir, "tables")
                        os.makedirs(tables_dir, exist_ok=True)
                        merged_filename = f"merged_table_{match.segment1.page_index}_{match.segment2.page_index}.png"
                        merged_path = os.path.join(tables_dir, merged_filename)
                        merged_img.save(merged_path)
                        
                        abs_merged_path = os.path.abspath(merged_path)
                        rel_merged = os.path.relpath(abs_merged_path, out_dir)
                        
                        pages_str = f"pages {match.segment1.page_index}-{match.segment2.page_index}"
                        
                        if self.vlm is not None:
                            wrote_table = False
                            try:
                                table = self.vlm.extract_table(abs_merged_path)
                                item = to_structured_dict(table)
                                if item:
                                    item["page"] = f"{match.segment1.page_index}-{match.segment2.page_index}"
                                    item["type"] = "Table (Merged)"
                                    item["split_merge"] = True
                                    item["merge_confidence"] = match.confidence
                                    structured_items.append(item)
                                    
                                    table_md = render_markdown_table(
                                        item.get("headers"), 
                                        item.get("rows"),
                                        title=item.get("title") or f"Merged Table ({pages_str})"
                                    )
                                    table_html = render_html_table(
                                        item.get("headers"), 
                                        item.get("rows"),
                                        title=item.get("title") or f"Merged Table ({pages_str})"
                                    )
                                    
                                    md_lines.append(f"\n### Merged Table ({pages_str})\n")
                                    md_lines.append(table_md)
                                    html_lines.append(f'<h3>Merged Table ({pages_str})</h3>')
                                    html_lines.append(table_html)
                                    wrote_table = True
                            except Exception as e:
                                pass
                            
                            if not wrote_table:
                                table_md = f"![Merged Table — {pages_str}]({rel_merged})\n"
                                table_html = f'<img src="{rel_merged}" alt="Merged Table — {pages_str}" />'
                                md_lines.append(f"\n### Merged Table ({pages_str})\n")
                                md_lines.append(table_md)
                                html_lines.append(f'<h3>Merged Table ({pages_str})</h3>')
                                html_lines.append(table_html)
                        else:
                            table_md = f"![Merged Table — {pages_str}]({rel_merged})\n"
                            table_html = f'<img src="{rel_merged}" alt="Merged Table — {pages_str}" />'
                            md_lines.append(f"\n### Merged Table ({pages_str})\n")
                            md_lines.append(table_md)
                            html_lines.append(f'<h3>Merged Table ({pages_str})</h3>')
                            html_lines.append(table_html)
                        
                        if tables_bar: tables_bar.update(1)
                        
                    except Exception as e:
                        print(f"⚠️  Warning: Failed to merge table {match_idx + 1}: {e}")

        md_path = write_markdown(md_lines, out_dir)
        
        if self.vlm is not None and html_lines:
            html_path = write_html_from_lines(html_lines, out_dir)
        else:
            html_path = write_html(md_lines, out_dir)
        
        pages_dir = os.path.join(out_dir, "pages")
        os.makedirs(pages_dir, exist_ok=True)
        
        for page_num, content_lines in page_content.items():
            page_md_path = os.path.join(pages_dir, f"page_{page_num:03d}.md")
            write_markdown(content_lines, os.path.dirname(page_md_path), os.path.basename(page_md_path))
        
        excel_path = None
        html_structured_path = None
        if self.vlm is not None and structured_items:
            excel_path = os.path.join(out_dir, "tables.xlsx")
            write_structured_excel(excel_path, structured_items)
            html_structured_path = os.path.join(out_dir, "tables.html")
            write_structured_html(html_structured_path, structured_items)

        print(f"✅ Enhanced parsing completed successfully!")
        print(f"📁 Output directory: {out_dir}")

    def _create_enhanced_pdf_from_pages(self, enhanced_pages: List[Image.Image], output_path: str) -> None:
        """
        Create an enhanced PDF from already processed enhanced pages.
        
        :param enhanced_pages: List of enhanced PIL images
        :param output_path: Path for the enhanced PDF
        """
        if not enhanced_pages:
            raise ValueError("No enhanced pages provided")
        
        try:
            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 from processed pages: {output_path}")
        except Exception as e:
            print(f"❌ Error creating enhanced PDF from pages: {e}")
            raise

    def restore_pdf_only(self, pdf_path: str, output_path: str = None, task: str = None) -> str:
        """
        Apply DocRes restoration to a PDF without parsing.
        
        :param pdf_path: Path to the input PDF file
        :param output_path: Path for the enhanced PDF (if None, auto-generates)
        :param task: DocRes restoration task (if None, uses instance default)
        :return: Path to the enhanced PDF or None if failed
        """
        if not self.use_image_restoration or not self.docres_engine:
            raise RuntimeError("Image restoration is not enabled or DocRes engine is not available")
        
        task = task or self.restoration_task
        return self.docres_engine.restore_pdf(pdf_path, output_path, task, self.restoration_dpi)

    def get_restoration_info(self) -> Dict[str, Any]:
        """
        Get information about the current restoration configuration.
        
        :return: Dictionary with restoration settings and status
        """
        return {
            'enabled': self.use_image_restoration,
            'task': self.restoration_task,
            'device': self.restoration_device,
            'dpi': self.restoration_dpi,
            'engine_available': self.docres_engine is not None,
            'supported_tasks': self.docres_engine.get_supported_tasks() if self.docres_engine else []
        }
