"""Visualization utilities for CVAT tools."""

import logging
import re
from pathlib import Path
from typing import Literal, Optional

from docling_core.types.doc.document import (
    ContentLayer,
    DoclingDocument,
    ImageRefMode,
    KeyValueItem,
)
from PIL import Image, ImageDraw, ImageFont

from docling_cvat_tools.utils import from_pil_to_base64
from docling_cvat_tools.visualisation.constants import HTML_DEFAULT_HEAD_FOR_COMP_v2

VizMode = Literal["reading_order", "key_value"]
ALL_CONTENT_LAYERS: set[ContentLayer] = {layer for layer in ContentLayer}


def _document_has_key_value_items(doc: DoclingDocument) -> bool:
    """Return True when the document includes at least one key-value annotation."""
    for item, _ in doc.iterate_items(
        traverse_pictures=True,
        included_content_layers=ALL_CONTENT_LAYERS,
    ):
        if isinstance(item, KeyValueItem):
            return True
    return False


def _visualization_filename_for_mode(base_filename: Path, viz_mode: VizMode) -> Path:
    """Derive the output filename for a visualization mode."""
    if viz_mode == "reading_order":
        return base_filename.with_name(
            f"{base_filename.stem}_layout{base_filename.suffix}"
        )
    if viz_mode == "key_value":
        return base_filename.with_name(
            f"{base_filename.stem}_key_value{base_filename.suffix}"
        )
    raise ValueError(f"Unsupported visualization mode: {viz_mode}")


def get_missing_pageimg(
    width: int = 800, height: int = 1100, text: str = "MISSING PAGE"
) -> Image.Image:
    """Get missing page image.

    Args:
        width: Image width
        height: Image height
        text: Text to display on the image

    Returns:
        PIL Image with the missing page text
    """
    # Create a white background image
    image = Image.new("RGB", (width, height), color="white")
    draw = ImageDraw.Draw(image)

    # Try to use a standard font or fall back to default
    try:
        # For larger installations, you might have Arial or other fonts
        font = ImageFont.truetype("arial.ttf", size=60)
    except IOError:
        # Fall back to default font
        font = ImageFont.load_default().font_variant(size=60)  # type: ignore

    # Get text size to center it
    text_width, text_height = (
        draw.textsize(text, font=font)
        if hasattr(draw, "textsize")
        else (draw.textlength(text, font=font), font.size)
    )

    # Position for the text (centered and angled)
    position = ((width - text_width) // 2, (height - text_height) // 2)

    # Draw the watermark text (light gray and rotated)
    draw.text(position, text, fill=(200, 200, 200), font=font)

    # Rotate the image 45 degrees to create diagonal watermark effect
    image = image.rotate(45, expand=False, fillcolor="white")

    return image


def _get_document_visualization_data(
    doc: DoclingDocument,
    page_no: int,
    pattern: re.Pattern[str],
    viz_mode: VizMode,
) -> tuple[str, str]:
    """Get visualization data for a document page.

    Args:
        doc: Document to visualize
        page_no: Page number to visualize
        pattern: Regex pattern to extract body content
        viz_mode: Visualization mode

    Returns:
        Tuple of (base64_image, html_content)
    """
    visualization_error: Optional[str] = None

    try:
        page_imgs = doc.get_visualization(show_label=False, viz_mode=viz_mode)
    except (IndexError, ValueError) as exc:
        visualization_error = str(exc)
        logging.error(
            "Could not render visualization for page %s due to geometry issue: %s",
            page_no,
            exc,
        )
        page_imgs = {}

    if page_no in page_imgs:
        doc_img_b64 = from_pil_to_base64(page_imgs[page_no])
    else:
        if visualization_error is None:
            logging.error(f"{page_no} not in page_imgs, get default image.")
        doc_img_b64 = from_pil_to_base64(get_missing_pageimg())

    try:
        doc_page = doc.export_to_html(
            image_mode=ImageRefMode.EMBEDDED,
            page_no=page_no,
            included_content_layers={ContentLayer.BODY, ContentLayer.FURNITURE},
        )
    except (IndexError, ValueError) as exc:
        logging.error(
            "Could not export page %s to HTML due to serialization error: %s",
            page_no,
            exc,
        )
        doc_page_body = (
            "<p>ERROR: Could not render page HTML due to serialization issues. "
            f"Details: {exc}</p>"
        )
        return doc_img_b64, doc_page_body

    # Search for the pattern in the HTML string
    mtch = pattern.search(doc_page)
    if mtch:
        doc_page_body = mtch.group(1).strip()
    else:
        logging.error(f"could not find body in doc_page")
        doc_page_body = "<p>Nothing Found</p>"

    if len(doc_page_body) == 0:
        doc_page_body = "<p>Nothing Found</p>"

    if visualization_error is not None:
        doc_page_body = (
            "<p>WARNING: Page visualization unavailable due to geometry issues. "
            f"Details: {visualization_error}</p>\n" + doc_page_body
        )

    return doc_img_b64, doc_page_body


def _create_visualization_html(
    filename: Path,
    true_doc: DoclingDocument,
    template: str,
    pred_doc: Optional[DoclingDocument] = None,
    draw_reading_order: bool = True,
    viz_mode: VizMode = "reading_order",
) -> None:
    """Create and save HTML visualization.

    Args:
        filename: Path to save the visualization
        true_doc: Document to visualize
        template: HTML template to use
        pred_doc: Optional predicted document for comparison
        draw_reading_order: Whether to draw reading order
        viz_mode: Visualization mode
    """
    # Compile regex pattern once
    pattern = re.compile(
        r"<body[^>]*>\n<div class='page'>(.*?)</div>\n</body>",
        re.DOTALL | re.IGNORECASE,
    )

    # Add CSS class based on view type
    view_class = "comparison-view" if pred_doc is not None else "single-view"
    template = template.replace(
        "<style>",
        f"<style>\n.{view_class} td {{ width: {'25%' if pred_doc is not None else '50%'}; }}",
    )

    html_parts = [
        "<!DOCTYPE html>",
        "<html>",
        template,
        "<body>",
    ]

    html_parts.append(f"<table class='{view_class}'>")
    html_parts.append("<tbody>")

    # Get page numbers and convert to set of integers
    true_page_nos = set(true_doc.pages.keys())
    if pred_doc is not None:
        pred_page_nos = set(pred_doc.pages.keys())
        if true_page_nos != pred_page_nos:
            logging.error(
                f"incompatible page numbers: \n"
                f"true pages: {true_page_nos}\npred pages: {pred_page_nos}"
            )
        page_nos = true_page_nos | pred_page_nos
    else:
        page_nos = true_page_nos

    # Process each page
    for page_no in page_nos:
        html_parts.append("<tr>")

        # Process true document
        true_img_b64, true_html = _get_document_visualization_data(
            true_doc,
            page_no,
            pattern,
            viz_mode,
        )
        html_parts.append("<td>")
        html_parts.append(f'<img src="data:image/png;base64,{true_img_b64}">')
        html_parts.append("</td>")
        html_parts.append("<td>")
        html_parts.append(f"<div class='page'>\n{true_html}\n</div>")
        html_parts.append("</td>")

        # Process predicted document if present
        if pred_doc is not None:
            pred_img_b64, pred_html = _get_document_visualization_data(
                pred_doc, page_no, pattern, viz_mode
            )
            html_parts.append("<td>")
            html_parts.append(f'<img src="data:image/png;base64,{pred_img_b64}">')
            html_parts.append("</td>")
            html_parts.append("<td>")
            html_parts.append(f"<div class='page'>\n{pred_html}\n</div>")
            html_parts.append("</td>")

        html_parts.append("</tr>")

    html_parts.append("</tbody>")
    html_parts.append("</table>")
    html_parts.extend(["</body>", "</html>"])

    with open(str(filename), "w") as fw:
        fw.write("\n".join(html_parts))


def save_single_document_html(
    filename: Path,
    doc: DoclingDocument,
    draw_reading_order: bool = True,
) -> None:
    """Save single document visualizations with their HTML content.

    Args:
        filename: Path used as the base visualization filename
        doc: Document to visualize
        draw_reading_order: Whether to draw reading order

    Always produces the default reading order visualization. The output is written
    to the provided base filename with an added ``_layout`` suffix. When the
    document holds key-value annotations, a second visualization is saved alongside
    it using a ``_key_value`` suffix.
    """
    viz_modes: list[VizMode] = ["reading_order"]
    if _document_has_key_value_items(doc):
        viz_modes.append("key_value")

    for viz_mode in viz_modes:
        viz_filename = _visualization_filename_for_mode(filename, viz_mode)
        _create_visualization_html(
            filename=viz_filename,
            true_doc=doc,
            template=HTML_DEFAULT_HEAD_FOR_COMP_v2,
            draw_reading_order=draw_reading_order,
            viz_mode=viz_mode,
        )
