"""
Pipe diagram generation utilities.

This module provides functionality to generate pipe diagrams from inspection data
using Cairo for rendering. It supports both NASSCO and WRC standards.
"""

import logging
import math
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import cairo

from ..config.pdf_config import PDFConfig
from ..models.types import MainlineInspectionUnion
from ..records.nassco.nassco_defects import NasscoDefectRecords
from ..records.nassco.survey_directions import SurveyDirection as NasscoSurveyDirection
from ..records.wrc.directions import Direction as WRCSurveyDirection
from ..records.wrc.wrc_defects import WrcDefectRecords
from ..utils.defect_properties import DefectBand
from ..utils.xml_models import XMLModel

logger = logging.getLogger(__name__)


@dataclass
class DiagramConfig:
    """Configuration for pipe diagram generation."""

    # Canvas dimensions
    width: int = 125
    height: int = 545
    height_multi_page: int = 600

    # Cairo surface size limits to prevent "invalid value (typically too big)" errors
    max_cairo_height: int = 30000  # Conservative limit to avoid Cairo surface size errors
    max_cairo_width: int = 30000  # Conservative limit to avoid Cairo surface size errors

    # Pipe and circle dimensions
    pipe_width: int = 10
    circle_radius: int = 20

    # Text positioning
    vertical_shift: float = 1.1
    horizontal_shift: float = 0.8
    text_padding: int = 10
    downstream_text_extra_gap: int = 6

    # Baseline Y start for first circle and pipe start
    # Increased to provide room for wrapped upstream manhole text above the circle
    start_y: int = 90

    # Colors
    background_color: Tuple[float, float, float, float] = (1.0, 1.0, 1.0, 1.0)
    pipe_color: Tuple[float, float, float, float] = (0.7, 0.7, 0.7, 0.65)
    border_color: Tuple[float, float, float] = (0.2, 0.2, 0.2)
    text_color: Tuple[float, float, float] = (0.2, 0.2, 0.2)

    # Font settings
    font_family: str = "Arial"
    font_size: int = 10
    observation_font_size: int = 8
    manhole_text_max_width: int = 80  # Max width for manhole text before wrapping
    manhole_text_line_spacing: int = 12  # Line spacing for wrapped manhole text

    # Manhole text positioning - extra margin to avoid observation line overlap
    upstream_text_top_margin: int = 50  # Extra margin above upstream circle for text
    downstream_text_bottom_margin: int = 25  # Extra margin below downstream circle for text

    # Pipe length calculation parameters
    pipe_top_margin: int = 200
    pipe_bottom_margin: int = 220

    # Observation line spacing
    observation_line_spacing: int = 36


class PipeDiagramGenerator:
    """
    Generates pipe diagrams from inspection data.

    This class handles the generation of pipe diagrams using Cairo for rendering.
    It supports both NASSCO (PACP) and WRC (Mainline) standards.

    Note: Manhole doesn't have pipe diagrams.
    """

    # Connection codes for different standards
    TAPFACTORY_NASSCO = [
        "TF",
        "TB",
        "TS",
        "TR",
        "TFA",
        "TBA",
        "TSA",
        "TRA",
        "TFD",
        "TBD",
        "TSD",
        "TRD",
        "TFC",
        "TBC",
        "TSC",
        "TRC",
        "TFI",
        "TBI",
        "TSI",
        "TRI",
        "TFB",
        "TBB",
        "TSB",
        "TRB",
    ]

    CONNECTION_WRC = ["CN"]

    # Survey abandoned codes
    SURVEY_ABANDONED_CODES = ["SA", "MSA"]

    DEFAULT_OBSERVATIONS_PER_PAGE = PDFConfig.default_observations_per_page

    def __init__(self, config: Optional[DiagramConfig] = None):
        """Initialize the pipe diagram generator."""
        self.config = config or DiagramConfig()
        self.surface: Optional[cairo.ImageSurface] = None
        self.ctx: Optional[cairo.Context] = None
        self.is_first_circle = True
        self.upstream_manhole = ""
        self.downstream_manhole = ""
        self.drawing_bounds = [float("inf"), float("inf"), 0, 0]  # [min_x, min_y, max_x, max_y]

    def _update_drawing_bounds(self, x: float, y: float, width: float = 0, height: float = 0):
        """Update the drawing bounds to track actual content area."""
        self.drawing_bounds[0] = min(self.drawing_bounds[0], x)
        self.drawing_bounds[1] = min(self.drawing_bounds[1], y)
        self.drawing_bounds[2] = max(self.drawing_bounds[2], x + width)
        self.drawing_bounds[3] = max(self.drawing_bounds[3], y + height)

    def _calculate_required_surface_height(self, end_y: float, downstream_manhole_id: str) -> int:
        """Calculate the required surface height to accommodate downstream manhole text."""
        if not downstream_manhole_id or downstream_manhole_id == "None":
            return self.config.height_multi_page

        # Calculate wrapped text height to account for multi-line text
        text_height = self._measure_wrapped_text_height(downstream_manhole_id, self.config.manhole_text_max_width, self.config.manhole_text_line_spacing)

        # Add extra padding for text descenders and bottom margin
        bottom_padding = 15
        text_y = end_y + self.config.circle_radius + self.config.downstream_text_bottom_margin
        required_height = int(text_y + text_height + bottom_padding)  # Add padding at bottom

        # Cap the height to prevent Cairo surface size errors
        # Cairo has limits around 32,767 pixels for surface dimensions

        # Return the maximum of required height and default multi-page height, but cap it
        calculated_height = max(required_height, self.config.height_multi_page)
        return min(calculated_height, self.config.max_cairo_height)

    def _create_context(self, surface: cairo.ImageSurface) -> cairo.Context:
        """Create and configure a Cairo context."""
        ctx = cairo.Context(surface)
        ctx.set_antialias(cairo.ANTIALIAS_BEST)
        ctx.set_source_rgba(*self.config.background_color)
        ctx.select_font_face(self.config.font_family, cairo.FONT_SLANT_NORMAL, cairo.FONT_WEIGHT_BOLD)
        ctx.set_font_size(self.config.font_size)
        ctx.rectangle(0, 0, self.config.width, self.config.height)
        ctx.fill()
        ctx.save()
        return ctx

    def _measure_text(self, text: str) -> Tuple[float, float]:
        """Measure text width and height using a temporary Cairo context."""
        temp_surface = cairo.ImageSurface(cairo.FORMAT_RGB24, 1, 1)
        temp_ctx = cairo.Context(temp_surface)
        temp_ctx.select_font_face(self.config.font_family, cairo.FONT_SLANT_NORMAL, cairo.FONT_WEIGHT_BOLD)
        temp_ctx.set_font_size(self.config.font_size)
        extents = temp_ctx.text_extents(text)
        return extents.width, extents.height

    def _measure_wrapped_text_height(self, text: str, max_width: float, line_spacing: float) -> float:
        """Measure the total height of wrapped text."""
        temp_surface = cairo.ImageSurface(cairo.FORMAT_RGB24, 1, 1)
        temp_ctx = cairo.Context(temp_surface)
        temp_ctx.select_font_face(self.config.font_family, cairo.FONT_SLANT_NORMAL, cairo.FONT_WEIGHT_BOLD)
        temp_ctx.set_font_size(self.config.font_size)

        # Wrap text using temporary context
        words = text.split()
        if not words:
            return line_spacing

        lines = []
        current_line = words[0]

        for word in words[1:]:
            test_line = current_line + " " + word
            test_extents = temp_ctx.text_extents(test_line)
            if test_extents.width <= max_width:
                current_line = test_line
            else:
                lines.append(current_line)
                # Check if word itself is too long, wrap by chars
                word_extents = temp_ctx.text_extents(word)
                if word_extents.width > max_width:
                    # Simple char wrapping for measurement
                    char_line = ""
                    for char in word:
                        test_char = char_line + char
                        if temp_ctx.text_extents(test_char).width > max_width and char_line:
                            lines.append(char_line)
                            char_line = char
                        else:
                            char_line = test_char
                    current_line = char_line
                else:
                    current_line = word

        if current_line:
            lines.append(current_line)

        num_lines = len(lines) if lines else 1
        return num_lines * line_spacing

    def _cap_pixel_for_end_text(self, pixel: float, total_distance: float, height_cap: int) -> float:
        """Cap pixel so end circle and downstream text fit within a height cap."""
        if total_distance <= 0:
            return pixel
        # Calculate wrapped text height to account for multi-line text
        text_h = self._measure_wrapped_text_height(self.downstream_manhole, self.config.manhole_text_max_width, self.config.manhole_text_line_spacing)
        # Use downstream_text_bottom_margin to match the actual positioning in _draw_circle
        # Add extra padding for text descenders and bottom margin
        bottom_padding = 15
        allowed_end_space = self.config.circle_radius + self.config.downstream_text_bottom_margin + text_h + bottom_padding
        max_pipe_end_y = height_cap - allowed_end_space
        if max_pipe_end_y <= self.config.start_y:
            return pixel
        pixel_cap = (max_pipe_end_y - self.config.start_y) / total_distance
        return min(pixel, pixel_cap)

    def _calculate_content_bounds(self, observations: List[Dict[str, Any]], standard: str, use_multi_page_height: bool = False) -> Tuple[int, int, int, int]:
        """Calculate the actual bounds of the content to be drawn."""
        if not observations:
            return (0, 0, 80, 200)  # Default size

        # Calculate the maximum Y position based on observations
        max_distance = max(obs["position"] for obs in observations)

        # Calculate the actual height needed based on the pipe length and positioning
        # Start position is configurable (Y coordinate for first circle)
        start_y = self.config.start_y

        # Use the appropriate height based on whether this is for multi-page or single-page
        available_height = self.config.height_multi_page if use_multi_page_height else self.config.height

        # Calculate the scale factor for the pipe length
        # This matches the calculation in _generate_single_page and _generate_multi_page
        if len(observations) < self.DEFAULT_OBSERVATIONS_PER_PAGE:
            h = ((available_height - self.config.pipe_top_margin) / self.DEFAULT_OBSERVATIONS_PER_PAGE) * len(observations)
            pixel = h / max_distance if max_distance > 0 else 1
        else:
            h = available_height - self.config.pipe_bottom_margin
            pixel = h / max_distance if max_distance > 0 else 1

        # Calculate the actual end Y position of the pipe
        pipe_end_y = start_y + (pixel * max_distance)

        # Add space for the end circle
        end_circle_y = pipe_end_y + self.config.circle_radius

        # Calculate the maximum Y position for observation lines
        # Each observation line extends to the right with increasing Y values
        max_observation_y = 50 + (len(observations) * self.config.observation_line_spacing)  # Based on _extend_line calculation

        # The total height needed is the maximum of pipe end + circle vs observation lines
        total_height = max(end_circle_y + 20, max_observation_y + 20)  # Add small padding

        # Width needed: pipe width + observation lines + text space
        # Observation lines extend to x + 25 + 20 = x + 45
        content_width = 80  # Enough for pipe + observation lines + text

        return (0, 0, content_width, int(total_height))

    def _crop_to_content(self, surface: cairo.ImageSurface, bounds: Tuple[int, int, int, int]) -> cairo.ImageSurface:
        """Crop the surface to just the content bounds."""
        x, y, width, height = bounds

        # Create a new surface with the content dimensions
        cropped_surface = cairo.ImageSurface(cairo.FORMAT_RGB24, width, height)
        cropped_ctx = cairo.Context(cropped_surface)

        # Fill with white background
        cropped_ctx.set_source_rgb(1.0, 1.0, 1.0)
        cropped_ctx.rectangle(0, 0, width, height)
        cropped_ctx.fill()

        # Copy the content from the original surface
        # The bounds represent the region to copy from the original surface
        cropped_ctx.set_source_surface(surface, -x, -y)
        cropped_ctx.paint()

        return cropped_surface

    def _get_severity_color(self, severity: int, standard: str) -> Tuple[float, float, float]:
        """Get color based on severity and standard."""
        if standard == "N":  # NASSCO
            return DefectBand(severity=severity).nassco_band_color
        else:  # WRC
            return DefectBand(severity=severity).wrc_band_color

    def _is_observation_code(self, code: str, standard: str) -> bool:
        """Check if a code is a valid observation code for the given standard."""
        if standard == "N":
            return code in [defect.value for defect in NasscoDefectRecords]
        else:
            return code in [defect.value for defect in WrcDefectRecords]

    def _show_text(self, x: float, y: float, radius: float, text: str) -> None:
        """Display text at the specified position."""
        if not self.ctx:
            return

        if text == self.upstream_manhole and text != "None":
            dist = y - (self.config.vertical_shift * self.config.circle_radius)
        elif text == "None":
            dist = y
        elif self._is_observation_code(text, "N") or self._is_observation_code(text, "W"):
            dist = y - 3.5
        elif text == self.downstream_manhole and text != "None":
            dist = y + ((self.config.vertical_shift + 0.3) * self.config.circle_radius)
        else:
            dist = y - 3.5

        text_x = x - (self.config.horizontal_shift * radius)

        # Update bounds for text
        text_extents = self.ctx.text_extents(text)
        self._update_drawing_bounds(text_x, dist, text_extents.width, text_extents.height)

        self.ctx.move_to(text_x, dist)
        self.ctx.show_text(text)

    def _wrap_text_by_chars(self, text: str, max_width: float) -> List[str]:
        """Wrap text by characters to fit within max_width."""
        if not self.ctx:
            return [text]

        lines = []
        current_line = ""

        for char in text:
            test_line = current_line + char
            test_extents = self.ctx.text_extents(test_line)
            if test_extents.width <= max_width:
                current_line = test_line
            else:
                if current_line:
                    lines.append(current_line)
                current_line = char

        if current_line:
            lines.append(current_line)

        return lines if lines else [text]

    def _wrap_text(self, text: str, max_width: float) -> List[str]:
        """Wrap text to fit within max_width, returning list of lines."""
        if not self.ctx:
            return [text]

        # Check if text fits in one line
        text_extents = self.ctx.text_extents(text)
        if text_extents.width <= max_width:
            return [text]

        # Split text into words and build lines
        words = text.split()

        # If no spaces in text, wrap by characters
        if len(words) <= 1:
            return self._wrap_text_by_chars(text, max_width)

        lines = []
        current_line = words[0]

        # Check if first word itself is too long
        first_word_extents = self.ctx.text_extents(current_line)
        if first_word_extents.width > max_width:
            # First word is too long, wrap it by characters
            char_wrapped = self._wrap_text_by_chars(current_line, max_width)
            lines.extend(char_wrapped[:-1])
            current_line = char_wrapped[-1] if char_wrapped else ""

        for word in words[1:]:
            test_line = current_line + " " + word
            test_extents = self.ctx.text_extents(test_line)
            if test_extents.width <= max_width:
                current_line = test_line
            else:
                if current_line:
                    lines.append(current_line)
                # Check if this word itself is too long
                word_extents = self.ctx.text_extents(word)
                if word_extents.width > max_width:
                    char_wrapped = self._wrap_text_by_chars(word, max_width)
                    lines.extend(char_wrapped[:-1])
                    current_line = char_wrapped[-1] if char_wrapped else ""
                else:
                    current_line = word

        if current_line:
            lines.append(current_line)

        return lines if lines else [text]

    def _draw_wrapped_text(self, x: float, y: float, text: str, max_width: float, line_spacing: float, align_center: bool = True, above: bool = True) -> float:
        """
        Draw text with wrapping support.

        Args:
            x: Center x position for the text
            y: Starting y position (top of text area if above=False, bottom if above=True)
            text: The text to draw
            max_width: Maximum width before wrapping
            line_spacing: Vertical spacing between lines
            align_center: If True, center each line horizontally at x
            above: If True, text extends upward from y; if False, downward

        Returns:
            The total height of the drawn text block
        """
        if not self.ctx:
            return 0

        lines = self._wrap_text(text, max_width)

        # Calculate total height
        total_height = len(lines) * line_spacing

        # Determine starting y position
        if above:
            # Text grows upward: start from bottom line
            start_y = y
            direction = -1
        else:
            # Text grows downward: start from top line
            start_y = y + line_spacing  # First line baseline
            direction = 1

        for i, line in enumerate(lines if not above else reversed(lines)):
            line_extents = self.ctx.text_extents(line)

            if align_center:
                text_x = x - line_extents.width / 2
            else:
                text_x = x

            text_y = start_y + (i * line_spacing * direction)

            # Update bounds for text
            self._update_drawing_bounds(text_x, text_y - line_extents.height, line_extents.width, line_extents.height)

            self.ctx.move_to(text_x, text_y)
            self.ctx.show_text(line)

        return total_height

    def _draw_circle(self, x: float, y: float, radius: float, grey_grad: float, manhole_id: str) -> None:
        """Draw a circle representing a manhole."""
        if not self.ctx:
            return
        # Update drawing bounds for the circle
        self._update_drawing_bounds(x - radius, y - radius, radius * 2, radius * 2)

        # Draw circle
        self.ctx.set_source_rgb(0.3, 0.3, 0.3)
        self.ctx.arc(x, y, radius, 0, 2 * math.pi)
        self.ctx.set_source_rgb(grey_grad, grey_grad, grey_grad)
        self.ctx.fill()

        # Draw border
        self.ctx.set_source_rgb(0.3, 0.3, 0.3)
        self.ctx.arc(x, y, radius, 0, 2 * math.pi)
        self.ctx.stroke()

        # Draw text outside circle with wrapping support
        # Position text with extra margin to avoid overlap with observation lines
        self.ctx.set_source_rgb(*self.config.text_color)
        max_width = self.config.manhole_text_max_width
        line_spacing = self.config.manhole_text_line_spacing

        if self.is_first_circle and manhole_id == self.upstream_manhole:
            # Upstream manhole: text at absolute top, well above observation lines
            text_y = y - radius - self.config.upstream_text_top_margin
            self._draw_wrapped_text(x, text_y, manhole_id, max_width, line_spacing, align_center=True, above=True)
            self.is_first_circle = False
        else:
            # Downstream manhole: text at absolute bottom, well below observation lines
            text_y = y + radius + self.config.downstream_text_bottom_margin
            self._draw_wrapped_text(x, text_y, manhole_id, max_width, line_spacing, align_center=True, above=False)

        self.ctx.move_to(x, y + radius)

    def _draw_rectangle(self, length: float, width: float, code: str) -> Tuple[float, float]:
        """Draw a rectangle representing a pipe segment."""
        if not self.ctx:
            return (0, 0)

        x, y = self.ctx.get_current_point()

        # Update drawing bounds for the rectangle
        self._update_drawing_bounds(x - width / 2, y, width, length)

        self.ctx.set_source_rgb(*self.config.border_color)
        stroke_ctx = cairo.Context(self.surface)
        stroke_ctx.set_line_width(1.2)
        stroke_ctx.set_source_rgb(*self.config.border_color)

        stroke_ctx.rectangle(x - (width / 2), y, width, length)
        stroke_ctx.stroke()

        self.ctx.set_source_rgba(*self.config.pipe_color)
        self.ctx.rectangle(x - (width / 2), y, width, length)
        self.ctx.fill()
        self.ctx.move_to(x, y + length)

        return self.ctx.get_current_point()

    def _draw_junction(self, pos: int) -> None:
        """Draw a junction/connection at the specified position."""
        if not self.ctx:
            return

        x, y = self.ctx.get_current_point()
        self.ctx.set_source_rgb(*self.config.border_color)

        if 1 <= pos <= 6:
            self.ctx.move_to(x - 1 * (self.config.pipe_width / 2) - 1, y - 5)
            self.ctx.rel_line_to(-7, -6)
            self.ctx.rel_line_to(-12, 0)
            self.ctx.rel_line_to(19, 18)
        elif 7 <= pos <= 12:
            self.ctx.move_to(x + 1 * (self.config.pipe_width / 2) + 1, y - 5)
            self.ctx.rel_line_to(7, -6)
            self.ctx.rel_line_to(12, 0)
            self.ctx.rel_line_to(-19, 18)

        self.ctx.close_path()
        self.ctx.set_source_rgba(*self.config.pipe_color)
        self.ctx.fill()
        self.ctx.move_to(x, y)

    def _draw_scb(self, radius: float, rect_length: float, rect_breadth: float) -> None:
        """Draw SCB (Special Connection Box) representation."""
        if not self.ctx:
            return

        x, y = self.ctx.get_current_point()

        self.ctx.save()
        self.ctx.translate(x, y)
        self.ctx.scale(1.5, 1)
        self.ctx.translate(-x, -y)

        self.ctx.set_source_rgb(*self.config.border_color)
        self.ctx.arc(x, y, radius / 2, 0, 2 * math.pi)
        self.ctx.fill()

        self.ctx.restore()
        self.ctx.move_to(x, y + radius / 2)

    def _extend_line(self, severity: int, code: str, index: int, standard: str, show_text: bool = True) -> None:
        """Extend a line to show observation details."""
        if not self.ctx:
            return

        color = self._get_severity_color(severity, standard)

        x, y = self.ctx.get_current_point()
        self.ctx.set_source_rgb(*color)
        new_y = ((index - 2) * self.config.observation_line_spacing) + 50

        # Update bounds for the observation line
        self._update_drawing_bounds(x, y, 25, new_y - y)  # Line from current point to new_y
        self._update_drawing_bounds(x + 25, new_y, 20, 0)  # Horizontal line

        self.ctx.move_to(x + self.config.circle_radius / 4, y)
        self.ctx.line_to(x + 25, new_y)

        a, b = self.ctx.get_current_point()
        self.ctx.set_line_width(1)
        self.ctx.rel_line_to(20, 0)
        self.ctx.stroke()

        self.ctx.set_font_size(self.config.observation_font_size)
        if show_text:
            self._show_text(a, b, 0, code)
        self.ctx.move_to(x, y)
        self.ctx.set_source_rgba(0.2, 0.2, 0.2, 0.9)

        # Ensure the text bounds are properly tracked - including the full text area
        text_extents = self.ctx.text_extents(code)
        text_y = b - text_extents.height  # Text baseline minus height
        self._update_drawing_bounds(a, text_y, text_extents.width, text_extents.height)

        # Also update bounds for the horizontal line endpoint
        self._update_drawing_bounds(a + 20, b, 0, 0)  # End of horizontal line

    def _abandoned_representation(self) -> None:
        """Draw representation for abandoned survey."""
        if not self.ctx:
            return

        self.ctx.set_source_rgb(0.9, 0, 0)
        self.ctx.rel_move_to(self.config.pipe_width / 2, 0)
        self.ctx.rel_move_to(10, -3)
        self.ctx.rel_line_to(-30, 9)
        self.ctx.rel_move_to(0, -4)
        self.ctx.rel_line_to(30, -9)
        self.ctx.stroke()

    def _save_image(self, file_number: int, output_dir: Path, observations: List[Dict[str, Any]], standard: str, is_final_page: bool = False) -> None:
        """Save the current surface as an image, cropped to content bounds."""
        if not self.surface:
            return

        # Use tracked drawing bounds if available, otherwise fall back to calculated bounds
        if self.drawing_bounds[0] != float("inf"):
            # Add some padding around the tracked bounds
            padding = 10

            # Use the correct height based on the surface dimensions
            surface_height = self.surface.get_height()
            surface_width = self.surface.get_width()

            # Calculate the actual content height needed
            content_height = int(self.drawing_bounds[3] - self.drawing_bounds[1] + 2 * padding)

            # If content extends beyond surface, we need to create a larger surface
            if content_height >= surface_height or self.drawing_bounds[3] >= surface_height:
                # Create a new surface with the required height
                new_surface = cairo.ImageSurface(cairo.FORMAT_RGB24, surface_width, content_height)
                new_ctx = cairo.Context(new_surface)
                new_ctx.set_source_rgb(1.0, 1.0, 1.0)  # White background
                new_ctx.rectangle(0, 0, surface_width, content_height)
                new_ctx.fill()

                # Copy the original surface to the new one
                new_ctx.set_source_surface(self.surface, 0, 0)
                new_ctx.paint()

                # Replace the original surface
                self.surface = new_surface
                self.ctx = new_ctx
                surface_height = content_height

            bounds = (
                max(0, int(self.drawing_bounds[0] - padding)),
                max(0, int(self.drawing_bounds[1] - padding)),
                min(surface_width, int(self.drawing_bounds[2] - self.drawing_bounds[0] + 2 * padding)),
                content_height,  # Use the content height, not surface height
            )
        else:
            # Fall back to calculated bounds - use multi-page height if surface is multi-page
            surface_height = self.surface.get_height()
            use_multi_page = surface_height == self.config.height_multi_page
            bounds = self._calculate_content_bounds(observations, standard, use_multi_page)

        cropped_surface = self._crop_to_content(self.surface, bounds)

        img_name = output_dir / f"pipe{file_number}.png"
        cropped_surface.write_to_png(str(img_name))
        cropped_surface.flush()

        # Only reset drawing bounds and create new surface if this is not the final page
        if not is_final_page:
            # Reset first circle flag for new page
            self.is_first_circle = True

            # Create new surface for next page with calculated height
            surface_height = getattr(self, "_required_height", self.config.height_multi_page)

            # Additional safety check to prevent Cairo surface size errors

            if surface_height > self.config.max_cairo_height or self.config.width > self.config.max_cairo_width:
                # Log warning and use safe fallback dimensions
                logger.warning(f"Surface dimensions too large for Cairo: width={self.config.width}, height={surface_height}. Using fallback dimensions to prevent Cairo error.")
                surface_height = min(surface_height, self.config.max_cairo_height)
                # Note: We don't modify width as it's part of the config, but we ensure height is safe

            self.surface = cairo.ImageSurface(cairo.FORMAT_RGB24, self.config.width, surface_height)
            self.ctx = self._create_context(self.surface)
            self.ctx.move_to(0, 0)
            self.ctx.set_source_rgba(*self.config.background_color)
            self.ctx.rectangle(0, 0, self.config.width, surface_height)
            self.ctx.fill()
            self.ctx.set_source_rgb(0, 0, 0)

            # Reset drawing bounds for the new page since it's a fresh surface
            self.drawing_bounds = [float("inf"), float("inf"), 0, 0]
        else:
            # For final page, clean up references to prevent potential memory leaks
            # (though this shouldn't be an issue since generator instances are not reused)
            self.surface = None
            self.ctx = None

    def _process_observations(self, observations: List[Any], standard: str) -> List[Dict[str, Any]]:
        """Process observations into a standardized format."""
        processed = []

        for obs in observations:
            if obs.distance is None:
                continue

            # Handle code extraction with error handling
            try:
                code = obs.code.value if obs.code else "None"
            except AttributeError:
                code = str(obs.code) if obs.code else "None"

            processed.append(
                {
                    "position": float(obs.distance),
                    "code": code,
                    "severity": obs.severity or 0,
                    "continuous_defect": obs.continuous_defect or "",
                    "clock_pos": obs.circumferential_location_from or 0,
                    "remarks": obs.remarks or "",
                }
            )

        # Sort by position
        processed.sort(key=lambda x: x["position"])
        return processed

    def _get_manhole_ids(self, inspection: XMLModel, standard: str) -> Tuple[str, str]:
        """Extract manhole IDs from inspection data."""
        try:
            if standard == "N":  # NASSCO standard
                upstream = str(inspection.header.measurements.upstream_mh_number or "").split(",")[-1]
                downstream = str(inspection.header.measurements.downstream_mh_number or "").split(",")[-1]
                direction = inspection.header.survey.direction_of_survey
            elif standard == "W":  # WRC standard
                upstream = str(inspection.header.measurements.start_node_reference or "").split(",")[-1]
                downstream = str(inspection.header.measurements.finish_node_reference or "").split(",")[-1]
                direction = inspection.header.survey.direction
            else:
                raise ValueError(f"Invalid standard: {standard}")

            # Handle direction
            if direction and direction in [NasscoSurveyDirection.upstream, WRCSurveyDirection.upstream]:
                upstream, downstream = downstream, upstream

        except AttributeError:
            # Fallback if field names are different
            upstream = "MH001"
            downstream = "MH002"

        return upstream, downstream

    def generate_diagram(self, inspection: XMLModel, output_dir: str | Path, single_page: bool = False, show_text: bool = True) -> List[Path]:
        """
        Generate pipe diagram from inspection data.

        Args:
            inspection: The inspection data (PACPObservation, MACPObservation, etc.)
            output_dir: Directory to save the generated images
            single_page: Whether to generate a single page diagram

        Returns:
            List of paths to generated image files
        """
        # Determine standard
        inspection_type = type(inspection).__name__
        if inspection_type in ["PACPInspection", "MACPInspection"]:
            standard = "N"
        else:
            standard = "W"

        # Get manhole IDs
        self.upstream_manhole, self.downstream_manhole = self._get_manhole_ids(inspection, standard)

        # Process observations
        observations = self._process_observations(inspection.observations, standard)

        if not observations:
            return []

        # Reset drawing bounds for new diagram
        self.drawing_bounds = [float("inf"), float("inf"), 0, 0]

        # Initialize surface
        self.surface = cairo.ImageSurface(cairo.FORMAT_RGB24, self.config.width, self.config.height)
        self.ctx = self._create_context(self.surface)

        if isinstance(output_dir, str):
            output_dir = Path(output_dir)

        # Generate diagram
        if single_page:
            return self._generate_single_page(observations, standard, output_dir, show_text)
        else:
            return self._generate_multi_page(observations, standard, output_dir, show_text)

    def _generate_single_page(self, observations: List[Dict[str, Any]], standard: str, output_dir: Path, show_text: bool = True) -> List[Path]:
        """Generate a single page diagram."""
        distances = [obs["position"] for obs in observations]
        total_observations = len(distances)

        h = self.config.height - 160  # Increased to account for starting at configurable Y
        # Check if last distance is zero to avoid division by zero
        last_distance = distances[total_observations - 1] if distances else 0
        pixel = h / last_distance if distances and last_distance > 0 else 1

        grey_grad = 0.8
        self._draw_circle(40, self.config.start_y, self.config.circle_radius, grey_grad, self.upstream_manhole)

        distance = 0
        no_of_obs = 1
        file_number = 0

        for row_num, obs in enumerate(observations):
            prev_distance = distance
            distance = obs["position"]
            length = abs(int((pixel * distance - pixel * prev_distance)))

            # Draw junction if needed
            if obs["code"] in self.TAPFACTORY_NASSCO or obs["code"] in self.CONNECTION_WRC:
                self._draw_junction(obs["clock_pos"])

            # Draw rectangle
            end_x, end_y = self._draw_rectangle(length, self.config.circle_radius / 2, obs["code"])

            # Handle special cases
            if obs["code"] == "GP" and "backdrop" in obs["remarks"].lower():
                self._extend_line(obs["severity"], obs["code"], row_num + 1, standard, show_text)
                self._draw_scb(self.config.circle_radius / 2, length, self.config.circle_radius / 2)
                end_x, end_y = self.ctx.get_current_point() if self.ctx else (0, 0)
            else:
                self._extend_line(obs["severity"], obs["code"], row_num + 1, standard, show_text)

            no_of_obs += 1

        # Draw end circle or abandoned representation
        if observations and observations[-1]["code"] in self.SURVEY_ABANDONED_CODES:
            self._abandoned_representation()
        else:
            if self.ctx:
                self.ctx.set_font_size(self.config.font_size)
                self._draw_circle(end_x, end_y + self.config.circle_radius, self.config.circle_radius, grey_grad, self.downstream_manhole)

        self._save_image(file_number, output_dir, observations, standard, is_final_page=True)
        return [output_dir / f"pipe{file_number}.png"]

    def _generate_multi_page(self, observations: List[Dict[str, Any]], standard: str, output_dir: Path, show_text: bool = True) -> List[Path]:
        """Generate a multi-page diagram."""
        distances = [obs["position"] for obs in observations]
        total_observations = len(distances)

        # Calculate pixel scale
        if total_observations < self.DEFAULT_OBSERVATIONS_PER_PAGE:
            h = ((self.config.height - self.config.pipe_top_margin) / self.DEFAULT_OBSERVATIONS_PER_PAGE) * total_observations
            # Check if last distance is zero to avoid division by zero
            last_distance = distances[total_observations - 1] if distances else 0
            pixel = h / last_distance if distances and last_distance > 0 else 1
            # Cap pixel to ensure the end circle + text fit within single-page height
            if distances and last_distance > 0:
                pixel = self._cap_pixel_for_end_text(pixel, last_distance, self.config.height)
        else:
            h = self.config.height - self.config.pipe_bottom_margin
            # Check if distance is zero to avoid division by zero
            ref_distance = distances[self.DEFAULT_OBSERVATIONS_PER_PAGE - 1] if len(distances) > self.DEFAULT_OBSERVATIONS_PER_PAGE - 1 else 0
            pixel = h / ref_distance if len(distances) > self.DEFAULT_OBSERVATIONS_PER_PAGE - 1 and ref_distance > 0 else 1

        # Calculate the required surface height for the downstream manhole
        # The downstream manhole will be at the end of the pipe
        last_distance = distances[total_observations - 1] if distances else 0
        pipe_end_y = self.config.start_y + (pixel * last_distance) if distances and last_distance > 0 else self.config.start_y
        required_height = self._calculate_required_surface_height(pipe_end_y, self.downstream_manhole)
        self._required_height = required_height  # Stored but not used to exceed caps

        grey_grad = 0.8
        self._draw_circle(40, self.config.start_y, self.config.circle_radius, grey_grad, self.upstream_manhole)

        length = 0
        highest_len = 0
        jump = 0
        obs_on_page = self.DEFAULT_OBSERVATIONS_PER_PAGE
        pg_no = 0
        distance = 0
        no_of_obs = 1
        file_number = 0
        generated_files = []
        page_observations = []  # Track observations for current page

        for row_num, obs in enumerate(observations):
            prev_distance = distance
            distance = obs["position"]
            length = abs(int((pixel * distance - pixel * prev_distance)))

            highest_len += length

            # Check if we need a new page BEFORE adding the observation
            if no_of_obs > obs_on_page:
                padding = 75
                remaining_obs = len(observations) - (no_of_obs - 1)

                if remaining_obs > PDFConfig.subsequent_page_observation_count:
                    distance_to_map = distances[min(row_num + (PDFConfig.subsequent_page_observation_count - 1), len(distances) - 1)] - distances[no_of_obs - 1]
                else:
                    distance_to_map = distances[total_observations - 1] - distances[no_of_obs - 1]
                    padding = 100

                if distance_to_map == 0:
                    distance_to_map = 1

                pixel = (self.config.height_multi_page - padding) / distance_to_map
                # If this becomes the last page, cap pixel so end circle + text fit within multi-page height
                if remaining_obs <= PDFConfig.subsequent_page_observation_count and distance_to_map > 0:
                    pixel = self._cap_pixel_for_end_text(pixel, distance_to_map, self.config.height_multi_page)
                highest_len = 0

                # Save the current page before starting a new one
                self._save_image(file_number, output_dir, page_observations, standard, is_final_page=False)
                generated_files.append(output_dir / f"pipe{file_number}.png")
                file_number += 1
                pg_no += 1

                if pg_no != 0:
                    obs_on_page = PDFConfig.subsequent_page_observation_count

                jump = 1
                no_of_obs = 1
                page_observations = []  # Reset for new page

            # Add observation to current page
            page_observations.append(obs)

            # Draw rectangle or handle jump
            if jump == 0:
                end_x, end_y = self._draw_rectangle(length, self.config.circle_radius / 2, obs["code"])
            else:
                if self.ctx:
                    self.ctx.move_to(40, 60)
                    end_x, end_y = self.ctx.get_current_point()
                jump = 0

            # Draw junction if needed
            if obs["code"] in self.TAPFACTORY_NASSCO or obs["code"] in self.CONNECTION_WRC:
                self._draw_junction(obs["clock_pos"])

            # Handle special cases
            if obs["code"] == "GP" and "backdrop" in obs["remarks"].lower():
                self._extend_line(obs["severity"], obs["code"], no_of_obs, standard, show_text)
                self._draw_scb(self.config.circle_radius / 2, length, self.config.circle_radius / 2)
                end_x, end_y = self.ctx.get_current_point() if self.ctx else (0, 0)
            else:
                self._extend_line(obs["severity"], obs["code"], no_of_obs, standard, show_text)

            no_of_obs += 1

        # Draw end circle or abandoned representation
        if observations and observations[-1]["code"] in self.SURVEY_ABANDONED_CODES:
            self._abandoned_representation()
        else:
            if self.ctx:
                self.ctx.set_font_size(self.config.font_size)
                self._draw_circle(end_x, end_y + self.config.circle_radius, self.config.circle_radius, grey_grad, self.downstream_manhole)

        # Save the final page if there are remaining observations
        if page_observations:
            self._save_image(file_number, output_dir, page_observations, standard, is_final_page=True)
            generated_files.append(output_dir / f"pipe{file_number}.png")

        return generated_files


def generate_pipe_diagrams(inspection: MainlineInspectionUnion, output_dir: str, show_text: bool = True):
    """
    Generate a pipe diagram for the Mainline.

    Args:
        inspection: The Mainline inspection to generate a pipe diagram for.
        output_dir: The path to the output directory.

    Returns:
        The paths to the generated pipe diagrams.
    """

    generator = PipeDiagramGenerator()
    files = generator.generate_diagram(inspection=inspection, output_dir=output_dir, single_page=False, show_text=show_text)
    return files
