"""
Word Document Preprocessor for DSOL V2 Extraction Pipeline.

This module preprocesses Word documents before Docling conversion by extracting
and embedding non-text content inline:

1. Images: Extract images → Generate AI captions → Embed as [IMAGE: caption]
2. Shapes: Extract text boxes/shapes → Embed as [SHAPE: text]
3. Comments: Extract comments → Embed as [COMMENT_1: author (date): text]

All content is embedded inline in the document text before Docling markdown conversion,
ensuring it's preserved in the semantic chunking phase.
"""

import base64
import tempfile
from pathlib import Path
from typing import List, Dict, Optional, Any

import structlog
from docx import Document
from docx.oxml.table import CT_Tbl
from docx.oxml.text.paragraph import CT_P
from docx.table import Table
from docx.text.paragraph import Paragraph

logger = structlog.get_logger(__name__)


class WordPreprocessor:
    """
    Preprocess Word documents before Docling conversion.

    Extraction and Inline Embedding:
    1. Images → AI-generated captions → Embed inline as [IMAGE: caption]
    2. Shapes/text boxes → Text extraction → Embed inline as [SHAPE: text]
    3. Comments → ID markers → Embed inline as [COMMENT_1: author (date): text]

    All content embedded inline before Docling markdown conversion.
    """

    def __init__(self, openai_client=None):
        """
        Initialize with Azure OpenAI client for image captioning.

        Args:
            openai_client: Optional Azure OpenAI client. If None, creates one lazily on first use.
        """
        self._client = openai_client
        self.logger = structlog.get_logger()

    def _get_client(self):
        """Get or create LLM client lazily (fork-safe)."""
        if self._client is None:
            from src.llm import get_chat_client

            self._client = get_chat_client()
            self.logger.info("Initialized LLM client for image captioning")

        return self._client

    def preprocess(self, file_path: Path) -> Path:
        """
        Preprocess Word document and save to new file.

        Processing steps:
        1. Load document
        2. Extract and caption images (AI-powered)
        3. Extract shapes and text boxes
        4. Extract and mark comments
        5. Save preprocessed document

        Args:
            file_path: Path to .docx file

        Returns:
            Path to preprocessed .docx file
        """
        self.logger.info("Starting Word preprocessing", file=file_path.name)

        # Load document
        doc = Document(str(file_path))

        # Step 1: Extract and caption images
        images = self._extract_images(doc)
        if images:
            self.logger.info("Extracting images", count=len(images))
            captions = self._generate_image_captions(images)
            self._embed_image_captions(doc, images, captions)
        else:
            self.logger.info("No images found in document")

        # Step 2: Extract shapes and text boxes
        shapes = self._extract_shapes(doc)
        if shapes:
            self.logger.info("Extracting shapes", count=len(shapes))
            self._embed_shape_descriptions(doc, shapes)
        else:
            self.logger.info("No shapes found in document")

        # Step 3: Extract and mark comments
        comments = self._extract_comments(doc)
        if comments:
            self.logger.info("Extracting comments", count=len(comments))
            self._apply_comment_markers(doc, comments)
        else:
            self.logger.info("No comments found in document")

        # Save preprocessed document
        output_path = file_path.parent / f"{file_path.stem}_preprocessed.docx"
        doc.save(str(output_path))

        self.logger.info(
            "Preprocessing complete",
            input=file_path.name,
            output=output_path.name,
            images=len(images),
            shapes=len(shapes),
            comments=len(comments)
        )

        return output_path

    def _extract_images(self, doc: Document) -> List[Dict]:
        """
        Extract all images from document with context.

        Returns list of dicts with:
        - rel_id: Relationship ID
        - data: Image binary data
        - context: Surrounding paragraphs text
        - paragraph: Paragraph containing the image
        - para_idx: Index of paragraph
        """
        images = []

        # Iterate through paragraphs to find images
        for para_idx, paragraph in enumerate(doc.paragraphs):
            # Check for images in paragraph runs
            for run in paragraph.runs:
                # Check for embedded images (drawings)
                if 'graphicData' in run._element.xml:
                    # Extract image relationship
                    blip_elements = run._element.xpath('.//a:blip')
                    for blip in blip_elements:
                        rel_id = blip.get('{http://schemas.openxmlformats.org/officeDocument/2006/relationships}embed')
                        if rel_id:
                            try:
                                # Get image data from relationship
                                image_part = doc.part.related_parts[rel_id]
                                img_data = image_part.blob

                                # Get context (previous and next 2 paragraphs)
                                context = self._get_paragraph_context(doc, para_idx, window=2)

                                images.append({
                                    'rel_id': rel_id,
                                    'data': img_data,
                                    'context': context,
                                    'paragraph': paragraph,
                                    'para_idx': para_idx
                                })
                            except KeyError:
                                self.logger.warning(
                                    "Image relationship not found",
                                    rel_id=rel_id,
                                    para_idx=para_idx
                                )

        return images

    def _get_paragraph_context(self, doc: Document, para_idx: int, window: int = 2) -> str:
        """
        Get text context around a paragraph.

        Args:
            doc: Word document
            para_idx: Index of target paragraph
            window: Number of paragraphs before and after to include

        Returns:
            Context string with surrounding paragraph text
        """
        start_idx = max(0, para_idx - window)
        end_idx = min(len(doc.paragraphs), para_idx + window + 1)

        context_paragraphs = []
        for i in range(start_idx, end_idx):
            text = doc.paragraphs[i].text.strip()
            if text:
                context_paragraphs.append(text)

        return " | ".join(context_paragraphs)

    def _generate_image_captions(self, images: List[Dict]) -> List[str]:
        """
        Generate AI captions for images using Azure OpenAI GPT-4o vision.

        Processes images individually with context from surrounding paragraphs.

        Args:
            images: List of image dictionaries with data and context

        Returns:
            List of caption strings (one per image)
        """
        captions = []

        for img in images:
            try:
                # Encode image to base64
                img_base64 = base64.b64encode(img['data']).decode('utf-8')

                # Determine image type from binary header
                img_type = self._detect_image_type(img['data'])

                # Call Azure OpenAI GPT-4o vision API
                client = self._get_client()
                from src.llm import get_chat_model_name

                response = client.chat.completions.create(
                    model=get_chat_model_name(),
                    messages=[
                        {
                            "role": "user",
                            "content": [
                                {
                                    "type": "text",
                                    "text": (
                                        f"Describe this image in 1-2 sentences. "
                                        f"Context from document: {img['context'][:200]}"
                                    )
                                },
                                {
                                    "type": "image_url",
                                    "image_url": {
                                        "url": f"data:image/{img_type};base64,{img_base64}"
                                    }
                                }
                            ]
                        }
                    ],
                    max_tokens=100
                )

                from src.llm import strip_thinking
                caption = strip_thinking(response.choices[0].message.content)
                captions.append(caption)

                self.logger.debug(
                    "Image caption generated",
                    para_idx=img['para_idx'],
                    caption=caption[:50]
                )

            except Exception as e:
                self.logger.warning(
                    "Image captioning failed",
                    para_idx=img.get('para_idx'),
                    error=str(e)
                )
                captions.append("[Image: Caption unavailable]")

        return captions

    def _detect_image_type(self, img_data: bytes) -> str:
        """
        Detect image type from binary header.

        Args:
            img_data: Image binary data

        Returns:
            Image type string (png, jpeg, gif, etc.)
        """
        if img_data.startswith(b'\x89PNG'):
            return "png"
        elif img_data.startswith(b'\xFF\xD8\xFF'):
            return "jpeg"
        elif img_data.startswith(b'GIF'):
            return "gif"
        elif img_data.startswith(b'BM'):
            return "bmp"
        else:
            # Default to png if unknown
            return "png"

    def _embed_image_captions(self, doc: Document, images: List[Dict], captions: List[str]):
        """
        Embed AI-generated captions inline near images in document.

        Adds caption text as inline marker after the image in the same paragraph.

        Args:
            doc: Word document
            images: List of image dictionaries
            captions: List of AI-generated captions (parallel to images)
        """
        for img, caption in zip(images, captions):
            para = img['paragraph']

            # Add caption as inline text marker after the image
            caption_text = f"[IMAGE: {caption}]"

            # Insert caption run in the same paragraph
            para.add_run(f" {caption_text}")

    def _extract_shapes(self, doc: Document) -> List[Dict]:
        """
        Extract shapes, text boxes, and drawing objects.

        Returns list of dicts with:
        - type: Shape type (textbox, shape, diagram)
        - text: Extracted text content
        - para_idx: Paragraph index
        """
        shapes = []

        # Shapes are embedded in paragraphs, similar to images
        for para_idx, paragraph in enumerate(doc.paragraphs):
            # Check for text boxes and shapes in runs
            for run in paragraph.runs:
                # Check for text box content
                if 'txbxContent' in run._element.xml:
                    # Extract text from text box
                    textbox_text = self._extract_textbox_text(run._element)
                    if textbox_text:
                        shapes.append({
                            'type': 'textbox',
                            'text': textbox_text,
                            'para_idx': para_idx,
                            'paragraph': paragraph
                        })

        return shapes

    def _extract_textbox_text(self, element) -> Optional[str]:
        """
        Extract text from a text box element.

        Args:
            element: XML element containing text box

        Returns:
            Extracted text or None if extraction fails
        """
        try:
            # Define namespace
            from lxml import etree
            ns = {'w': 'http://schemas.openxmlformats.org/wordprocessingml/2006/main'}

            # Navigate XML structure to find text
            textbox_paras = element.findall('.//w:p', ns)
            texts = []
            for para_elem in textbox_paras:
                text_elems = para_elem.findall('.//w:t', ns)
                para_text = ''.join(t.text for t in text_elems if t.text)
                if para_text:
                    texts.append(para_text)
            return ' | '.join(texts)
        except Exception as e:
            self.logger.debug(f"Failed to extract textbox text: {e}")
            return None

    def _embed_shape_descriptions(self, doc: Document, shapes: List[Dict]):
        """
        Embed shape text inline into document flow.

        Adds shape content as inline marker at the shape's location.

        Args:
            doc: Word document
            shapes: List of shape dictionaries
        """
        for shape in shapes:
            # Find the paragraph and append shape text as inline marker
            if 'paragraph' in shape and shape['paragraph']:
                para = shape['paragraph']
                para.add_run(f" [SHAPE: {shape['text']}]")

    def _extract_comments(self, doc: Document) -> List[Dict]:
        """
        Extract all comments from document.

        Returns list of dicts with:
        - id: Comment ID
        - author: Comment author
        - text: Comment text
        - date: Comment date
        """
        comments = []

        try:
            # Access comments part
            if hasattr(doc.part, 'comments_part') and doc.part.comments_part:
                for comment_elem in doc.part.comments_part.element:
                    comment_id = comment_elem.get('{http://schemas.openxmlformats.org/wordprocessingml/2006/main}id')
                    author = comment_elem.get('{http://schemas.openxmlformats.org/wordprocessingml/2006/main}author')
                    date = comment_elem.get('{http://schemas.openxmlformats.org/wordprocessingml/2006/main}date')

                    # Extract comment text
                    text_elems = comment_elem.xpath('.//w:t', namespaces={
                        'w': 'http://schemas.openxmlformats.org/wordprocessingml/2006/main'
                    })
                    comment_text = ''.join(t.text for t in text_elems if t.text)

                    if comment_text:
                        comments.append({
                            'id': comment_id,
                            'author': author or 'Unknown',
                            'text': comment_text,
                            'date': date or 'N/A'
                        })
        except AttributeError:
            # No comments part exists
            self.logger.debug("No comments part found in document")

        return comments

    def _apply_comment_markers(self, doc: Document, comments: List[Dict]):
        """
        Apply comment markers inline in document.

        Embeds comments at the end of the document as inline markers for semantic chunking.

        Args:
            doc: Word document
            comments: List of comment dictionaries
        """
        if not comments:
            return

        # Add comments section at end (will be included in markdown conversion)
        doc.add_paragraph()  # Blank line
        doc.add_paragraph("DOCUMENT COMMENTS:")

        for comment in comments:
            comment_para = doc.add_paragraph()
            comment_para.add_run(
                f"[COMMENT_{comment['id']}] {comment['author']} ({comment['date']}): {comment['text']}"
            )