"""
PDF tool handlers for MCP server
Implements list_pdfs and read_pdf tools with intelligent limits

Backend: pdfplumber (text/tables) + pypdfium2 (page rendering) + pikepdf (image extraction)
"""
import pdfplumber
import pikepdf
import pypdfium2 as pdfium
from PIL import Image
import io
import logging
import os
import re
from pathlib import Path
from typing import Any

from src.config import config
from .schemas import (
    ListPDFsInput, ListPDFsOutput, PDFInfo, ListPDFsError,
    ReadPDFInput, ReadPDFSuccess, ReadPDFError,
    PageData, ContentBlock
)
from .validators import validate_pdf_read_request
from .image_processor import crop_image_to_max_dimension, is_header_footer_image
from .file_matcher import find_similar_pdfs, get_file_not_found_message
from .text_validator import is_text_corrupted, check_pdf_corruption_with_pdfminer
from .table_converter import convert_table_to_markdown
from .content_orderer import order_content_blocks, merge_adjacent_text_blocks
from .text_extractor import extract_non_table_text_regions

logger = logging.getLogger(__name__)


def validate_path_security(path: Path, allowed_base: Path) -> bool:
    """
    Validate that a path is within the allowed base directory.
    Prevents path traversal attacks.

    Args:
        path: Path to validate
        allowed_base: Base directory that path must be within

    Returns:
        True if path is safe, False if potential path traversal
    """
    try:
        # Resolve both paths to eliminate symlinks and .. components
        resolved_path = path.resolve()
        resolved_base = allowed_base.resolve()

        # Check if the resolved path starts with the resolved base
        return resolved_path.is_relative_to(resolved_base)
    except (ValueError, RuntimeError):
        return False


async def list_pdfs_handler(arguments: dict[str, Any]) -> tuple[str, list[dict]]:
    """
    List all PDF files recursively from working directory

    Args:
        arguments: Dictionary matching ListPDFsInput schema

    Returns:
        Tuple of (JSON string, empty list) - no images for list_pdfs
    """
    try:
        # Validate input
        input_data = ListPDFsInput(**arguments)

        # Resolve working directory
        working_dir = Path(input_data.working_directory)

        # If not absolute path, use current directory as base
        if not working_dir.is_absolute():
            working_dir = Path.cwd() / working_dir

        working_dir = working_dir.resolve()

        if not working_dir.exists():
            # List all subdirectories of the current directory
            current_dir = Path.cwd()
            subdirs = []

            # Direct subdirectories of the current directory
            try:
                for item in sorted(current_dir.iterdir()):
                    if item.is_dir() and not item.name.startswith('.'):
                        subdirs.append(item.name)
            except (PermissionError, OSError) as e:
                logger.debug(f"Could not list directory {current_dir}: {e}")

            if subdirs:
                dir_list = "\n".join([f"  - {d}" for d in subdirs])
                msg = (
                    f"Working directory not found: {working_dir}\n\n"
                    f"Available folders in current directory ({current_dir}):\n{dir_list}"
                )
            else:
                msg = f"Working directory not found: {working_dir}"

            error = ListPDFsError(error="DIRECTORY_NOT_FOUND", message=msg)
            return error.model_dump_json(indent=2), []

        if not working_dir.is_dir():
            error = ListPDFsError(
                error="NOT_A_DIRECTORY",
                message=f"Path is not a directory: {working_dir}"
            )
            return error.model_dump_json(indent=2), []

        # Find PDF files
        pdfs = []
        pattern = "**/*.pdf" if input_data.include_subdirectories else "*.pdf"
        max_depth = input_data.max_depth

        for pdf_path in sorted(working_dir.glob(pattern)):
            try:
                # Skip if not a file
                if not pdf_path.is_file():
                    continue

                # Check depth limit
                if max_depth is not None and input_data.include_subdirectories:
                    relative_path = pdf_path.relative_to(working_dir)
                    depth = len(relative_path.parts) - 1  # -1 for filename
                    if depth > max_depth:
                        continue

                # Get page count using pypdfium2
                pdf_doc = pdfium.PdfDocument(str(pdf_path))
                total_pages = len(pdf_doc)
                pdf_doc.close()

                # Add to list
                pdfs.append(PDFInfo(
                    name=pdf_path.name,
                    path=str(pdf_path),
                    pages=total_pages,
                    size_bytes=pdf_path.stat().st_size
                ))

            except Exception:
                # Skip files that can't be opened (corrupted, not PDF, etc.)
                continue

        # Create output
        output = ListPDFsOutput(
            pdfs=pdfs,
            total_count=len(pdfs),
            working_directory=str(working_dir)
        )

        return output.model_dump_json(indent=2), []

    except PermissionError as e:
        error = ListPDFsError(
            error="PERMISSION_DENIED",
            message=f"Permission denied: {str(e)}"
        )
        return error.model_dump_json(indent=2), []
    except Exception as e:
        # Return error using schema
        error = ListPDFsError(
            error="INTERNAL_ERROR",
            message=f"Failed to list PDFs: {str(e)}"
        )
        return error.model_dump_json(indent=2), []


async def read_pdf_handler(arguments: dict[str, Any]) -> tuple[str, list[dict]]:
    """
    Read PDF with page and image limits, intelligent validation

    Implements:
    1. Pre-scan validation (page/image limits)
    2. Text/image ordering preservation (from comclerk_agent)
    3. Header/footer filtering
    4. Image cropping to A4 1/4 (from qwen-code)
    5. Suggested ranges on limit exceeded

    Args:
        arguments: Dictionary matching ReadPDFInput schema

    Returns:
        Tuple of (JSON string, list of image dicts)
        - JSON has image placeholders like [IMAGE_0], [IMAGE_1]
        - Images list: [{"data": bytes, "format": "jpeg"|"png"}, ...]
          (MCP SDK Image wrapper handles base64 encoding)
    """
    try:
        # Validate input
        input_data = ReadPDFInput(**arguments)

        # Resolve PDF path
        pdf_path = Path(input_data.file_path)
        if not pdf_path.is_absolute():
            pdf_path = Path.cwd() / pdf_path
        pdf_path = pdf_path.resolve()

        # Check file exists
        if not pdf_path.exists():
            # Find similar files
            similar_files = find_similar_pdfs(str(pdf_path), max_suggestions=3)
            error_message = get_file_not_found_message(str(pdf_path), similar_files)

            error = ReadPDFError(
                error="FILE_NOT_FOUND",
                message=error_message,
                suggested_files=similar_files if similar_files else None
            )
            return error.model_dump_json(indent=2), []

        # Security: Check for path traversal attacks
        # The file must be within the current working directory or its subdirectories
        cwd = Path.cwd().resolve()
        if not validate_path_security(pdf_path, cwd):
            error = ReadPDFError(
                error="PERMISSION_DENIED",
                message="Access denied: File path must be within the current working directory"
            )
            return error.model_dump_json(indent=2), []

        # Check read permission
        if not os.access(pdf_path, os.R_OK):
            error = ReadPDFError(
                error="PERMISSION_DENIED",
                message=f"Permission denied: Cannot read file {pdf_path}"
            )
            return error.model_dump_json(indent=2), []

        # Validate limits BEFORE processing
        validation = validate_pdf_read_request(
            str(pdf_path),
            input_data.start_page,
            input_data.end_page
        )

        if not validation.valid:
            # Return error with suggested ranges
            error = ReadPDFError(
                error=validation.error,
                message=validation.message,
                total_pages=validation.total_pages,
                total_images=validation.total_images,
                suggested_ranges=validation.suggested_ranges
            )
            return error.model_dump_json(indent=2), []

        # Open PDF with pdfplumber for text extraction
        # Also open pikepdf once for image extraction (if needed)
        with pdfplumber.open(str(pdf_path)) as pdf:
            total_pages = len(pdf.pages)
            end_page = min(input_data.end_page or total_pages, total_pages)

            pages_data = []
            total_images_count = 0
            extracted_images = []  # List of {"data": bytes, "format": "jpeg"|"png"}
            image_index = 0  # Global image index for placeholders

            # Check if we need pikepdf for image extraction
            mode = input_data.extraction_mode
            need_pikepdf = (mode == 'auto')  # Only auto mode extracts document images

            # Open pikepdf once outside the loop (if needed)
            pike_pdf = None
            if need_pikepdf:
                try:
                    pike_pdf = pikepdf.open(str(pdf_path))
                except Exception as e:
                    logger.warning(f"Failed to open PDF with pikepdf: {e}")

            try:
                # Process pages in order (preserves text/image ordering)
                for page_num in range(input_data.start_page, end_page + 1):
                    if page_num > total_pages:
                        break

                    page = pdf.pages[page_num - 1]  # 0-based index

                    # Check extraction mode (mode already set above)
                    skip_text_extraction = (mode == 'image_only')

                    # Variables for text extraction hint (image_only mode)
                    extractable_char_count = None
                    text_hint = None

                    if skip_text_extraction:
                        # Skip text and table extraction entirely
                        tables_with_position = []
                        table_bboxes = []
                        text_lines_for_ordering = []
                        full_text = ""

                        # Check for extractable text (to provide hint)
                        page_text = page.extract_text() or ''
                        char_count = len(page_text.strip())
                        if char_count > 0:
                            extractable_char_count = char_count

                            # Check corruption to determine hint
                            pdfminer_corrupted, warning_count = check_pdf_corruption_with_pdfminer(
                                str(pdf_path), page_num
                            )
                            if pdfminer_corrupted:
                                is_corrupted = True
                                ratio = warning_count / 10
                            else:
                                is_corrupted, ratio = is_text_corrupted(page_text)

                            if is_corrupted:
                                corruption_pct = int(ratio * 100)
                                text_hint = f"{char_count} chars ({corruption_pct}% corrupted). Text extraction not recommended."
                            else:
                                text_hint = f"{char_count} chars extractable. Use 'auto' to get text."
                    else:
                        # Extract tables first (to get bboxes)
                        tables_raw = page.extract_tables()
                        tables_with_position = []
                        table_bboxes = []

                        if tables_raw:
                            table_objects = page.find_tables()
                            for i, (table_data, table_obj) in enumerate(zip(tables_raw, table_objects)):
                                md = convert_table_to_markdown(table_data) if table_data else ""
                                if md:
                                    tables_with_position.append({
                                        'top': table_obj.bbox[1],
                                        'markdown': f"**Table {i+1}**\n\n{md}"
                                    })
                                    table_bboxes.append(table_obj.bbox)

                        # Extract text EXCLUDING table regions (prevent duplication!)
                        text_regions = extract_non_table_text_regions(page, table_bboxes)

                        # Use text_regions instead of text_lines
                        text_lines_for_ordering = []
                        for region in text_regions:
                            # Treat each region as a single "line"
                            text_lines_for_ordering.append({
                                'top': region['top'],
                                'text': region['text']
                            })

                        full_text = '\n\n'.join(r['text'] for r in text_regions)  # Full text (excluding tables)

                    # Auto-detect text corruption (only if extracting text)
                    text_corrupted = False
                    corruption_ratio = 0.0

                    if not skip_text_extraction:
                        pdfminer_corrupted, warning_count = check_pdf_corruption_with_pdfminer(
                            str(pdf_path), page_num
                        )

                        # Use pdfminer warnings first, fallback to character-based detection
                        if pdfminer_corrupted:
                            text_corrupted = True
                            corruption_ratio = warning_count / 10
                        else:
                            # Include table content in corruption check
                            all_text_for_check = full_text
                            if tables_with_position:
                                table_texts = '\n\n'.join(t['markdown'] for t in tables_with_position)
                                all_text_for_check = f"{full_text}\n\n{table_texts}" if full_text else table_texts

                            text_corrupted, corruption_ratio = is_text_corrupted(all_text_for_check)

                    # Provide full page as image based on extraction mode
                    page_image_base64 = None
                    page_img_width = None
                    page_img_height = None

                    # Determine if we should include page image
                    should_include_image = (
                        mode == 'image_only' or                  # Always in image_only mode
                        (mode == 'auto' and text_corrupted)      # Auto mode: only if corrupted
                    )
                    # mode == 'text_only' never includes page image

                    if should_include_image:
                        try:
                            logger.info(f"Page {page_num}: Attempting to render page image (mode={mode}, text_corrupted={text_corrupted})")
                            # Render page image with pypdfium2 (no external dependencies!)
                            pdfium_doc = pdfium.PdfDocument(str(pdf_path))
                            pdfium_page = pdfium_doc[page_num - 1]  # 0-indexed

                            # Calculate scale based on DPI (PDF base is 72 DPI)
                            scale = input_data.page_image_dpi / 72
                            bitmap = pdfium_page.render(scale=scale)
                            page_img = bitmap.to_pil()
                            pdfium_doc.close()

                            logger.info(f"Page {page_num}: pypdfium2 rendered page as {page_img.size}")

                            if page_img:

                                # White-out header/footer if filter enabled (for image_only mode)
                                if input_data.filter_header_footer:
                                    from PIL import ImageDraw
                                    draw = ImageDraw.Draw(page_img)
                                    width, height = page_img.size
                                    # Header: top 6% of page
                                    header_height = int(height * config.header_footer_ratio)
                                    draw.rectangle([0, 0, width, header_height], fill='white')
                                    # Footer: bottom 6% of page
                                    footer_start = int(height * config.footer_start_ratio)
                                    draw.rectangle([0, footer_start, width, height], fill='white')

                                # Resize page image to fit within max_image_dimension
                                width, height = page_img.size
                                max_dim = input_data.max_image_dimension
                                if width > max_dim or height > max_dim:
                                    scale = min(max_dim / width, max_dim / height)
                                    new_width = int(width * scale)
                                    new_height = int(height * scale)
                                    page_img = page_img.resize((new_width, new_height), Image.Resampling.LANCZOS)

                                # Convert to RGB if necessary (JPEG doesn't support RGBA)
                                if page_img.mode in ('RGBA', 'LA', 'P'):
                                    page_img = page_img.convert('RGB')

                                buffered = io.BytesIO()
                                page_img.save(buffered, format="JPEG", quality=config.jpeg_quality)

                                # Add to extracted images with raw bytes (MCP SDK handles base64 encoding)
                                extracted_images.append({
                                    "data": buffered.getvalue(),  # raw bytes
                                    "format": "jpeg"              # format for Image wrapper
                                })
                                logger.info(f"Page {page_num}: Added page image to extracted_images, total={len(extracted_images)}")
                                page_image_base64 = f"[IMAGE_{image_index}]"
                                image_index += 1
                                total_images_count += 1  # Count page images too

                                page_img_width = page_img.width
                                page_img_height = page_img.height

                                # Note: text_hint not set here since text_corrupted
                                # already triggers corruption warning in UI
                            else:
                                logger.warning(f"Page {page_num}: pypdfium2 failed to render page")
                        except Exception as e:
                            import traceback
                            logger.error(f"Page {page_num}: Failed to render page image: {e}")
                            logger.error(f"Page {page_num}: Traceback: {traceback.format_exc()}")

                    # Extract images with pikepdf (skip if image_only or text_only mode)
                    images_with_position = []

                    # text_only mode: skip all image extraction
                    skip_image_extraction = (mode == 'image_only' or mode == 'text_only')

                    if not skip_image_extraction and pike_pdf:
                        # Get image position info from pdfplumber
                        pdfplumber_images = page.images if hasattr(page, 'images') else []

                        # Extract actual image data with pikepdf (use pre-opened pike_pdf)
                        try:
                            pike_page = pike_pdf.pages[page_num - 1]

                            for img_index, (img_name, raw_image) in enumerate(pike_page.images.items()):
                                try:
                                    # Extract image using pikepdf
                                    pdf_image = pikepdf.PdfImage(raw_image)
                                    pil_image = pdf_image.as_pil_image()

                                    # Get display size from pdfplumber (actual size in PDF page)
                                    display_width = pil_image.width
                                    display_height = pil_image.height

                                    if img_index < len(pdfplumber_images):
                                        pdf_img_info = pdfplumber_images[img_index]
                                        # Calculate display size from bbox (in PDF points)
                                        # pdfplumber provides width/height in PDF points
                                        if 'width' in pdf_img_info and 'height' in pdf_img_info:
                                            # PDF points to pixels (assuming 72 DPI base, scale to reasonable size)
                                            # Limit to actual display size in PDF
                                            display_width = int(pdf_img_info['width'])
                                            display_height = int(pdf_img_info['height'])

                                    # Resize image to fit within display size while maintaining aspect ratio
                                    if pil_image.width > display_width or pil_image.height > display_height:
                                        from PIL import Image as PILImage
                                        # Calculate scale to fit within display size while keeping aspect ratio
                                        scale = min(display_width / pil_image.width, display_height / pil_image.height)
                                        new_width = int(pil_image.width * scale)
                                        new_height = int(pil_image.height * scale)
                                        pil_image = pil_image.resize(
                                            (new_width, new_height),
                                            PILImage.Resampling.LANCZOS
                                        )

                                    # Convert PIL image to bytes
                                    img_buffer = io.BytesIO()
                                    pil_image.save(img_buffer, format='PNG')
                                    img_bytes = img_buffer.getvalue()

                                    # Create img_data dict for header/footer filter
                                    img_data = {
                                        "image": img_bytes,
                                        "width": pil_image.width,
                                        "height": pil_image.height
                                    }

                                    # Filter header/footer if enabled
                                    if input_data.filter_header_footer:
                                        if is_header_footer_image(img_data):
                                            continue

                                    # Crop image if enabled (additional max dimension limit)
                                    if input_data.crop_images:
                                        result = crop_image_to_max_dimension(
                                            img_bytes,
                                            input_data.max_image_dimension
                                        )

                                        # Skip if image is too small and discarded
                                        if result[0] is None:
                                            continue

                                        img_bytes = result[0]

                                    # Add to extracted images with raw bytes (MCP SDK handles base64 encoding)
                                    extracted_images.append({
                                        "data": img_bytes,   # raw bytes
                                        "format": "png"      # format for Image wrapper
                                    })
                                    img_placeholder = f"[IMAGE_{image_index}]"
                                    image_index += 1

                                    # Get position from pdfplumber if available
                                    top = pdfplumber_images[img_index].get('top', 0) if img_index < len(pdfplumber_images) else 0

                                    images_with_position.append({
                                        'top': top,
                                        'image_data': img_placeholder  # Use placeholder instead of base64
                                    })
                                    total_images_count += 1

                                except Exception as e:
                                    logger.debug(f"Failed to extract image {img_index} on page {page_num}: {e}")
                                    continue
                        except Exception as e:
                            logger.warning(f"Failed to extract images with pikepdf on page {page_num}: {e}")

                    # Reconstruct content order (using text excluding tables)
                    ordered_blocks_raw = order_content_blocks(
                        text_lines_for_ordering,
                        tables_with_position,
                        images_with_position
                    )

                    # Merge adjacent text blocks
                    ordered_blocks = merge_adjacent_text_blocks(ordered_blocks_raw)

                    # Convert to ContentBlock (simple conversion!)
                    content_blocks = []

                    # CRITICAL: Skip adding corrupted text/tables if text is corrupted
                    # This applies to both Auto mode (text_corrupted=True) and image_only mode
                    should_skip_text_blocks = (
                        (mode == 'auto' and text_corrupted) or  # Auto mode: skip if corrupted
                        mode == 'image_only'                     # image_only: always skip
                    )

                    if should_skip_text_blocks:
                        # Only add image blocks (if any), skip text and table blocks
                        for block in ordered_blocks:
                            if block['type'] == 'image':
                                content_blocks.append(ContentBlock(
                                    type=block['type'],
                                    content=block.get('content')
                                ))
                    else:
                        # Normal mode: add all blocks
                        for block in ordered_blocks:
                            content = block.get('content')
                            # Clean up text blocks
                            if block['type'] == 'text' and content:
                                content = content.strip()
                                # Skip page number separators like "- 2 -", "- 10 -"
                                if re.match(r'^-\s*\d+\s*-$', content):
                                    continue
                            content_blocks.append(ContentBlock(
                                type=block['type'],
                                content=content
                            ))

                    pages_data.append(PageData(
                        page_number=page_num,
                        content_blocks=content_blocks,
                        text_corrupted=text_corrupted if text_corrupted else None,
                        corruption_ratio=corruption_ratio if text_corrupted else None,
                        page_image=page_image_base64,
                        page_image_width=page_img_width,
                        page_image_height=page_img_height,
                        extractable_char_count=extractable_char_count,
                        text_hint=text_hint
                    ))
            finally:
                # Close pikepdf if it was opened
                if pike_pdf:
                    pike_pdf.close()

        # Create success response
        output = ReadPDFSuccess(
            file_path=str(pdf_path),
            pages=pages_data,
            total_pages_read=len(pages_data),
            total_images=total_images_count
        )

        logger.info(f"Returning {len(extracted_images)} images")
        # Remove null values for clean output
        return output.model_dump_json(indent=2, exclude_none=True), extracted_images

    except FileNotFoundError as e:
        error = ReadPDFError(
            error="FILE_NOT_FOUND",
            message=f"PDF file not found: {str(e)}"
        )
        return error.model_dump_json(indent=2), []
    except PermissionError as e:
        error = ReadPDFError(
            error="PERMISSION_DENIED",
            message=f"Permission denied accessing PDF: {str(e)}"
        )
        return error.model_dump_json(indent=2), []
    except (pikepdf.PdfError, pdfplumber.pdfminer.pdfparser.PDFSyntaxError) as e:
        error = ReadPDFError(
            error="INVALID_PDF",
            message=f"Invalid or corrupted PDF file: {str(e)}"
        )
        return error.model_dump_json(indent=2), []
    except Exception as e:
        # Return error with more specific context
        error = ReadPDFError(
            error="INVALID_PDF",
            message=f"Error processing PDF: {type(e).__name__}: {str(e)}"
        )
        return error.model_dump_json(indent=2), []