# Copyright (c) 2025 CoReason, Inc.
#
# This software is proprietary and dual-licensed.
# Licensed under the Prosperity Public License 3.0 (the "License").
# A copy of the license is available at https://prosperitylicense.com/versions/3.0.0
# For details, see the LICENSE file.
# Commercial use beyond a 30-day trial requires a separate license.
#
# Source Code: https://github.com/CoReason-AI/coreason_auditor

import html
from typing import Any, List, Tuple

from reportlab.lib import colors
from reportlab.lib.pagesizes import letter
from reportlab.lib.styles import ParagraphStyle, getSampleStyleSheet
from reportlab.lib.units import inch
from reportlab.pdfgen import canvas
from reportlab.platypus import (
    Flowable,
    PageBreak,
    Paragraph,
    SimpleDocTemplate,
    Spacer,
    Table,
    TableStyle,
)

from coreason_auditor.models import AuditPackage, EventType, Session
from coreason_auditor.utils.logger import logger


class HorizontalLine(Flowable):  # type: ignore[misc]
    """Draws a horizontal line."""

    def __init__(self, width: float) -> None:
        super().__init__()
        self.width = width

    def wrap(self, availWidth: float, availHeight: float) -> Tuple[float, float]:
        return (self.width, 1)

    def draw(self) -> None:
        self.canv.setStrokeColor(colors.lightgrey)
        self.canv.line(0, 0, self.width, 0)


class PDFReportGenerator:
    """Generates the CoReason Audit Package PDF report."""

    def generate_report(self, audit_package: AuditPackage, output_path: str) -> None:
        """Generates a PDF report from the AuditPackage.

        Args:
            audit_package: The data to report.
            output_path: The file path to save the PDF.
        """
        logger.info(f"Generating PDF report for Audit Package {audit_package.id} at {output_path}")

        doc = SimpleDocTemplate(output_path, pagesize=letter)

        # Define the page template with header/footer
        def header_footer(canvas: canvas.Canvas, doc: Any) -> None:
            canvas.saveState()

            # --- Header ---
            canvas.setFont("Helvetica-Bold", 10)
            canvas.drawString(inch, letter[1] - 0.5 * inch, "CoReason Audit Report")

            canvas.setFont("Helvetica", 9)
            generated_str = f"Generated: {audit_package.generated_at.strftime('%Y-%m-%d %H:%M:%S UTC')}"
            canvas.drawRightString(letter[0] - inch, letter[1] - 0.5 * inch, generated_str)

            # Line separator
            canvas.setStrokeColor(colors.grey)
            canvas.line(inch, letter[1] - 0.6 * inch, letter[0] - inch, letter[1] - 0.6 * inch)

            # --- Footer ---
            canvas.setFont("Helvetica", 8)
            page_num = canvas.getPageNumber()
            canvas.drawString(inch, 0.5 * inch, "Confidential - CoReason Ecosystem")
            canvas.drawCentredString(letter[0] / 2, 0.5 * inch, f"Page {page_num}")

            if audit_package.document_hash:
                # Small hash on bottom right
                hash_short = (
                    audit_package.document_hash[:16] + "..."
                    if len(audit_package.document_hash) > 16
                    else audit_package.document_hash
                )
                canvas.drawRightString(letter[0] - inch, 0.5 * inch, f"Hash: {hash_short}")

            canvas.restoreState()

        # Build the story
        styles = getSampleStyleSheet()
        story: List[Any] = []

        # Custom Styles
        title_style = styles["Title"]
        heading_style = styles["Heading2"]
        normal_style = styles["Normal"]

        # Title Page Info
        story.append(Paragraph("CoReason Audit Package", title_style))
        story.append(Spacer(1, 24))

        header_data = [
            ["Generated At:", audit_package.generated_at.strftime("%Y-%m-%d %H:%M:%S UTC")],
            ["Agent Version:", audit_package.agent_version],
            ["Package ID:", str(audit_package.id)],
            ["Generated By:", audit_package.generated_by],
        ]
        story.append(self._create_info_table(header_data))
        story.append(Spacer(1, 24))

        # --- AI-BOM Section ---
        story.append(Paragraph("1. AI-BOM Inventory", heading_style))
        story.append(Spacer(1, 6))

        bom = audit_package.bom
        bom_data = [
            ["Model Identity:", Paragraph(html.escape(bom.model_identity), normal_style)],
        ]
        story.append(self._create_info_table(bom_data))
        story.append(Spacer(1, 12))

        # Data Lineage
        story.append(Paragraph("Data Lineage (Ingestion Jobs):", styles["Heading3"]))
        if bom.data_lineage:
            lineage_data = [[job] for job in bom.data_lineage]
            story.append(self._create_list_table(lineage_data, ["Job ID"]))
        else:
            story.append(Paragraph("No data lineage records found.", normal_style))
        story.append(Spacer(1, 12))

        # Dependencies
        story.append(Paragraph("Software Dependencies:", styles["Heading3"]))
        if bom.software_dependencies:
            dep_data = [[dep] for dep in bom.software_dependencies]
            story.append(self._create_list_table(dep_data, ["PackageSpec"]))
        else:
            story.append(Paragraph("No dependencies listed.", normal_style))
        story.append(Spacer(1, 24))

        # --- Traceability Matrix Section ---
        story.append(Paragraph("2. Requirement Traceability Matrix (RTM)", heading_style))
        story.append(Paragraph(f"Overall Status: {audit_package.rtm.overall_status.value}", normal_style))
        story.append(Spacer(1, 6))

        rtm_table_data = self._build_rtm_table_data(audit_package)
        story.append(self._create_data_table(rtm_table_data, col_widths=[60, 200, 80, 150]))
        story.append(Spacer(1, 24))

        # --- Deviation Report Section ---
        story.append(Paragraph("3. Deviation Report", heading_style))
        story.append(Spacer(1, 6))

        if audit_package.deviation_report:
            deviation_table_data = self._build_deviation_table_data(audit_package.deviation_report)
            story.append(self._create_data_table(deviation_table_data, col_widths=[100, 100, 60, 200]))
            story.append(Spacer(1, 24))

            # --- Detailed Transcripts ---
            # Added as requested to provide full forensic context
            story.append(Paragraph("4. Detailed Session Transcripts", heading_style))
            story.append(Spacer(1, 12))

            self._append_detailed_transcripts(story, audit_package.deviation_report)
        else:
            story.append(Paragraph("No deviations reported.", normal_style))

        story.append(Spacer(1, 36))

        # --- Config Changes Section ---
        story.append(Paragraph("5. Configuration Change Log", heading_style))
        story.append(Spacer(1, 6))

        if audit_package.config_changes:
            config_table_data = self._build_config_changes_table_data(audit_package.config_changes)
            # Adjusted widths to fit 468pt (Total ~450pt for safety margin)
            story.append(self._create_data_table(config_table_data, col_widths=[65, 45, 60, 60, 60, 100, 60]))
        else:
            story.append(Paragraph("No configuration changes recorded.", normal_style))
        story.append(Spacer(1, 24))

        # --- Signature Page ---
        self._append_signature_page(story, audit_package)

        # Build doc with header/footer callback
        doc.build(story, onFirstPage=header_footer, onLaterPages=header_footer)
        logger.info("PDF generation successful.")

    def _append_detailed_transcripts(self, story: List[Any], sessions: List[Session]) -> None:
        """
        Appends detailed transcripts for each session in the deviation report.
        Handles formatting for INPUT, THOUGHT, TOOL, and OUTPUT events.
        """
        styles = getSampleStyleSheet()

        # Define styles for different event types
        input_style = ParagraphStyle(
            "EventInput",
            parent=styles["Normal"],
            fontName="Helvetica-Bold",
            spaceBefore=6,
            spaceAfter=6,
            backColor=colors.whitesmoke,
            borderPadding=4,
        )

        thought_style = ParagraphStyle(
            "EventThought",
            parent=styles["Italic"],
            textColor=colors.dimgrey,
            spaceBefore=4,
            spaceAfter=4,
            leftIndent=12,
        )

        tool_style = ParagraphStyle(
            "EventTool",
            parent=styles["Code"],
            fontName="Courier",
            fontSize=8,
            textColor=colors.darkblue,
            spaceBefore=4,
            spaceAfter=4,
            leftIndent=12,
            backColor=colors.lightyellow,
            borderPadding=2,
        )

        output_style = ParagraphStyle(
            "EventOutput",
            parent=styles["Normal"],
            spaceBefore=6,
            spaceAfter=6,
            leftIndent=0,
        )

        header_style = styles["Heading3"]

        for session in sessions:
            # Session Header
            ts_str = session.timestamp.strftime("%Y-%m-%d %H:%M:%S UTC")
            header_text = f"Session: {session.session_id} | Time: {ts_str} | Risk: {session.risk_level.value}"
            story.append(Paragraph(header_text, header_style))

            # Violation Summary
            if session.violation_summary:
                violation_text = f"<b>Violation:</b> {html.escape(session.violation_summary)}"
                story.append(Paragraph(violation_text, styles["Normal"]))

            story.append(Spacer(1, 12))

            if not session.events:
                story.append(Paragraph("<i>No events recorded for this session.</i>", styles["Normal"]))

            for event in session.events:
                content_safe = html.escape(event.content).replace("\n", "<br/>")

                if event.event_type == EventType.INPUT:
                    label = "<b>User:</b> "
                    story.append(Paragraph(label + content_safe, input_style))

                elif event.event_type == EventType.THOUGHT:
                    label = "<b>Reasoning:</b> "
                    story.append(Paragraph(label + content_safe, thought_style))

                elif event.event_type == EventType.TOOL:
                    label = "<b>Tool Use:</b> "
                    story.append(Paragraph(label + content_safe, tool_style))

                elif event.event_type == EventType.OUTPUT:
                    label = "<b>Agent:</b> "
                    story.append(Paragraph(label + content_safe, output_style))

            # Separator between sessions
            story.append(Spacer(1, 24))
            story.append(HorizontalLine(width=450))
            story.append(Spacer(1, 24))

    def _append_signature_page(self, story: List[Any], audit_package: AuditPackage) -> None:
        """Appends the electronic signature page."""
        styles = getSampleStyleSheet()

        # Force page break so signature is on its own page (or clear separate section)
        story.append(PageBreak())

        story.append(Paragraph("Electronic Signature Page", styles["Heading1"]))
        story.append(Spacer(1, 24))

        story.append(Paragraph("This document has been electronically signed and approved.", styles["Normal"]))
        story.append(Spacer(1, 36))

        # Signature Block
        sig_content = [
            f"<b>Signed By:</b> {html.escape(audit_package.generated_by)}",
            f"<b>Date:</b> {audit_package.generated_at.strftime('%Y-%m-%d %H:%M:%S UTC')}",
            "<b>Reason:</b> Regulatory Compliance Submission",
            "<b>Authority:</b> CoReason Identity Service",
        ]

        if audit_package.electronic_signature:
            sig_content.append(f"<b>Signature Hash:</b> {audit_package.electronic_signature}")

        for line in sig_content:
            story.append(Paragraph(line, styles["Normal"]))
            story.append(Spacer(1, 6))

        story.append(Spacer(1, 24))
        story.append(Paragraph("<i>End of Document</i>", styles["Italic"]))

    def _create_info_table(self, data: List[List[Any]]) -> Table:
        """Creates a simple 2-column info table."""
        t = Table(data, colWidths=[120, 350])
        t.setStyle(
            TableStyle(
                [
                    ("ALIGN", (0, 0), (-1, -1), "LEFT"),
                    ("VALIGN", (0, 0), (-1, -1), "TOP"),
                    ("FONTNAME", (0, 0), (0, -1), "Helvetica-Bold"),
                    ("FONTNAME", (1, 0), (1, -1), "Helvetica"),
                    ("BOTTOMPADDING", (0, 0), (-1, -1), 6),
                ]
            )
        )
        return t

    def _create_list_table(self, data: List[List[str]], headers: List[str]) -> Table:
        """Creates a simple list table with a header."""
        table_data = [headers] + data
        t = Table(table_data, colWidths=[400], repeatRows=1)  # Repeat header if splits
        t.setStyle(
            TableStyle(
                [
                    ("BACKGROUND", (0, 0), (-1, 0), colors.lightgrey),
                    ("ALIGN", (0, 0), (-1, -1), "LEFT"),
                    ("GRID", (0, 0), (-1, -1), 0.5, colors.grey),
                    ("FONTNAME", (0, 0), (-1, 0), "Helvetica-Bold"),
                ]
            )
        )
        return t

    def _create_data_table(self, data: List[List[Any]], col_widths: List[int]) -> Table:
        """Creates a detailed data table with headers."""
        t = Table(data, colWidths=col_widths, repeatRows=1)
        t.setStyle(
            TableStyle(
                [
                    ("BACKGROUND", (0, 0), (-1, 0), colors.lightgrey),
                    ("ALIGN", (0, 0), (-1, -1), "LEFT"),
                    ("VALIGN", (0, 0), (-1, -1), "TOP"),
                    ("GRID", (0, 0), (-1, -1), 0.5, colors.black),
                    ("FONTNAME", (0, 0), (-1, 0), "Helvetica-Bold"),
                    ("FONTSIZE", (0, 0), (-1, -1), 9),
                    ("LEFTPADDING", (0, 0), (-1, -1), 4),
                    ("RIGHTPADDING", (0, 0), (-1, -1), 4),
                    ("TOPPADDING", (0, 0), (-1, -1), 4),
                    ("BOTTOMPADDING", (0, 0), (-1, -1), 4),
                ]
            )
        )
        return t

    def _build_rtm_table_data(self, audit_package: AuditPackage) -> List[List[Any]]:
        """Constructs the RTM table data rows."""
        headers = ["Req ID", "Description", "Status", "Tests"]
        rows = [headers]

        styles = getSampleStyleSheet()
        normal_style = styles["Normal"]

        # Map tests for lookup
        test_map = {t.test_id: t for t in audit_package.rtm.tests}

        for req in audit_package.rtm.requirements:
            req_id = req.req_id
            # Sanitize description to prevent XML parsing errors in Paragraph
            safe_desc = html.escape(req.desc)
            desc = Paragraph(safe_desc, normal_style)  # Wrap text

            linked_test_ids = audit_package.rtm.coverage_map.get(req_id, [])

            test_summaries = []
            req_status = "PASSED"

            if not linked_test_ids:
                req_status = "UNCOVERED"

            for tid in linked_test_ids:
                test = test_map.get(tid)
                if test:
                    test_summaries.append(f"{tid}: {test.result}")
                    if test.result != "PASS":
                        req_status = "FAILED"
                else:
                    test_summaries.append(f"{tid}: MISSING")
                    req_status = "FAILED"

            test_column_text = Paragraph(", ".join(test_summaries), normal_style)

            rows.append([req_id, desc, req_status, test_column_text])

        return rows

    def _build_deviation_table_data(self, deviations: List[Any]) -> List[List[Any]]:
        """Constructs the Deviation Report table data rows."""
        headers = ["Session ID", "Timestamp", "Risk", "Violation"]
        rows = [headers]

        styles = getSampleStyleSheet()
        normal_style = styles["Normal"]

        for session in deviations:
            sess_id = session.session_id
            # Format timestamp nicely
            ts = session.timestamp.strftime("%Y-%m-%d %H:%M:%S")
            risk = session.risk_level.value

            # Construct Violation Text (Type + Summary)
            summary_parts = []
            if session.violation_type:
                summary_parts.append(f"<b>{html.escape(session.violation_type)}</b>")

            if session.violation_summary:
                summary_parts.append(html.escape(session.violation_summary))

            if not summary_parts:
                violation_text = "No details"
            else:
                violation_text = ": ".join(summary_parts) if len(summary_parts) > 1 else summary_parts[0]

            violation_para = Paragraph(violation_text, normal_style)

            rows.append([sess_id, ts, risk, violation_para])

        return rows

    def _build_config_changes_table_data(self, changes: List[Any]) -> List[List[Any]]:
        """Constructs the Configuration Change Log table data rows."""
        headers = ["Date", "User", "Field", "From", "To", "Reason", "Status"]
        rows = [headers]

        styles = getSampleStyleSheet()
        normal_style = styles["Normal"]

        for change in changes:
            ts = change.timestamp.strftime("%Y-%m-%d")
            user = html.escape(change.user_id)
            field = html.escape(change.field_changed)

            # Wrap possibly long values
            old_val = Paragraph(html.escape(change.old_value), normal_style)
            new_val = Paragraph(html.escape(change.new_value), normal_style)
            reason = Paragraph(html.escape(change.reason), normal_style)

            status = html.escape(change.status)

            rows.append([ts, user, field, old_val, new_val, reason, status])

        return rows
