import os
from pathlib import Path
from typing import Optional

from typing_extensions import Literal

from ...config.pdf_config import PDFConfig
from ...models.wrc.mainline.inspection import MainlineInspection
from ...records.wrc.inspection_methods import InspectionMethod
from ...records.wrc.inspection_purposes import InspectionPurpose
from ...utils.pdf_reports import generate_pdf
from ...utils.pipe_diagram_generator import generate_pipe_diagrams
from ...utils.pipe_ratings import calculate_pipe_ratings


def generate_wrc_mainline_pdf(
    inspection: MainlineInspection,
    pipe_diagrams_path: str,
    annotated_frames_dir: str,
    pdf_output_path: str,
    video_name: str,
    uom: Literal["metric", "imperial"],
    logo_path: Optional[str] = None,
    direction_img: Optional[str] = None,
    address: Optional[str] = None,
    pdf_title: Optional[str] = None,
):
    """
    Generate a PDF report for the Mainline.

    Args:
        inspection: The Mainline inspection to generate a PDF report for.
        pipe_diagrams_path: The path to store the pipe diagrams.
        annotated_frames_dir: The path of the directory to store the annotated frames.
        pdf_output_path: The path to the output PDF file.
        video_name: The name of the video.
        uom: The unit of measurement.
        logo_path: The path to the logo image.
        direction_img: The path to the direction image.
        address: The address for the report.
        pdf_title: The title for the PDF report.
    Returns:
        The path to the generated PDF file.
    """

    inspection_dict = inspection.__dict__.copy()
    inspection_dict["observations"].sort(key=lambda x: x.distance)

    # Process purpose_of_inspection, method_of_inspection - use label for enums, string as-is
    # Convert nested objects to dictionaries for easier manipulation
    if "header" in inspection_dict and hasattr(inspection_dict["header"], "__dict__"):
        section_string = inspection_dict["header"].section_string
        inspection_dict["header"] = inspection_dict["header"].__dict__.copy()
        inspection_dict["header"]["section_string"] = section_string

    if "survey" in inspection_dict.get("header", {}) and hasattr(inspection_dict["header"]["survey"], "__dict__"):
        inspection_dict["header"]["survey"] = inspection_dict["header"]["survey"].__dict__.copy()

    if hasattr(inspection.header.survey, "purpose_of_inspection") and inspection.header.survey.purpose_of_inspection:
        try:
            purpose_of_inspection = InspectionPurpose.from_abbreviation(inspection.header.survey.purpose_of_inspection).label
        except ValueError:
            purpose_of_inspection = inspection.header.survey.purpose_of_inspection
        inspection_dict["header"]["survey"]["purpose_of_inspection"] = purpose_of_inspection

    if hasattr(inspection.header.survey, "method_of_inspection") and inspection.header.survey.method_of_inspection:
        try:
            method_of_inspection = InspectionMethod.from_abbreviation(inspection.header.survey.method_of_inspection).label
        except ValueError:
            method_of_inspection = inspection.header.survey.method_of_inspection
        inspection_dict["header"]["survey"]["method_of_inspection"] = method_of_inspection

    # Compute page-wise observation batching: first 12, then 18 per page
    total_observations = len(inspection_dict.get("observations", []))
    first_page_count = PDFConfig.first_page_observation_count
    subsequent_page_count = PDFConfig.subsequent_page_observation_count

    # Build counts for each page based on total observations
    page_counts = []
    if total_observations > 0:
        first = min(first_page_count, total_observations)
        page_counts.append(first)
        remaining = total_observations - first
        while remaining > 0:
            take = min(subsequent_page_count, remaining)
            page_counts.append(take)
            remaining -= take

    # Build per-page mapping of diagram image to observation slice
    diagram_pages = []
    offset = 0
    diagram_paths = generate_pipe_diagrams(inspection, pipe_diagrams_path, show_text=False)
    for page_index, path in enumerate(diagram_paths):
        count_for_page = page_counts[page_index] if page_index < len(page_counts) else 0
        page_observations = inspection_dict["observations"][offset : offset + count_for_page] if count_for_page > 0 else []
        diagram_pages.append(
            {
                "page_index": page_index,
                "diagram_path": path.absolute().as_uri(),
                "start_index": offset,
                "end_index": (offset + len(page_observations) - 1) if page_observations else -1,
                "observations": page_observations,
            }
        )
        offset += count_for_page

    pipe_ratings = calculate_pipe_ratings(inspection, uom.lower() == "imperial")

    # Enrich inspection with annotated frame URIs
    for index, observation in enumerate(inspection.observations):
        observation_dict = observation.__dict__.copy()
        annotated_frame_path = os.path.join(annotated_frames_dir, observation.image_reference)
        if not os.path.exists(annotated_frame_path):
            raise FileNotFoundError(f"Annotated frame file not found: {annotated_frame_path}")

        observation_dict["annotated_frame"] = Path(annotated_frame_path).absolute().as_uri()
        inspection_dict["observations"][index] = observation_dict

    measurement_units = "m" if uom == "metric" else "ft"
    pipe_size_units = "in" if uom == "imperial" else "mm"

    pipeline_length_reference = inspection.header.pipe.pipeline_length_reference if inspection.header.pipe.pipeline_length_reference else ""
    survey_date = inspection.header.survey.survey_date.strftime("%m%d%y") if inspection.header.survey.survey_date else ""
    image_prefix = ""
    if pipeline_length_reference:
        image_prefix = f"{pipeline_length_reference}-"
    if survey_date:
        image_prefix = f"{image_prefix}{survey_date}-"

    data = {
        "title": pdf_title or "Mainline Report",
        "inspection": inspection_dict,
        "diagram_pages": diagram_pages,
        "logo_path": logo_path,
        "video_name": video_name,
        "measurement_units": measurement_units,
        "pipe_size_units": pipe_size_units,
        "pipe_ratings": pipe_ratings,
        "direction_img": direction_img,
        "address": address,
        "image_prefix": image_prefix,
    }

    pdf_path = generate_pdf("wrc/mainline/report.html", data, pdf_output_path)
    return pdf_path
