import sys
import gc
import re
import json
import logging
import tempfile
from pathlib import Path
from datetime import datetime
import pdfplumber
from pypdf import PdfReader, PdfWriter
from pypdf.constants import UserAccessPermissions
from pypdf.generic import NameObject, BooleanObject
from typing import BinaryIO, TypedDict, Literal
from flask import Flask, request, send_file, abort
from werkzeug.utils import secure_filename

# --- Logging Configuration ---
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [%(levelname)s] %(name)s: %(message)s',
    handlers=[logging.StreamHandler(sys.stdout)]
)
logger = logging.getLogger(__name__)

app = Flask(__name__)

# Target template name inside the container working directory
TEMPLATE_NAME = "pftr_template.pdf"

InterpretationType = Literal[
    "normal",
    "pvd",
    "cw_neuro_disorder",
    "ild_pn",
    "asthma_cb",
    "emphysema"
]


class ParsedPdfData(TypedDict):
    patient_name: str
    date_of_service: str
    date_of_birth: str
    smoker: str
    highlighted_box: InterpretationType


def process_pft_pdf(input_stream: BinaryIO) -> BinaryIO:
    text_mapping: dict[str, InterpretationType] = {
        "normal": "normal",
        "pv disorders": "pvd",
        "pulmonary vascular": "pvd",
        "cw and nm disorders": "cw_neuro_disorder",
        "chest wall": "cw_neuro_disorder",
        "neuromuscular": "cw_neuro_disorder",
        "ild pneumonitis": "ild_pn",
        "interstitial lung disorders": "ild_pn",
        "pneumonitis": "ild_pn",
        "asthma cb": "asthma_cb",
        "asthma": "asthma_cb",
        "chronic bronchitis": "asthma_cb",
        "emphysema": "emphysema",
    }

    record: ParsedPdfData = {
        "patient_name": "",
        "date_of_service": "",
        "date_of_birth": "",
        "smoker": "",
        "highlighted_box": "normal"
    }

    input_stream.seek(0)

    # Use the upload stream directly so each request does not duplicate the PDF
    # in memory before parsing.
    with pdfplumber.open(input_stream) as pdf:
        if len(pdf.pages) < 2:
            logger.error(
                "Expected a 2-page PFT report, found %s page(s).", len(pdf.pages))
            raise ValueError(
                f"Expected a 2-page PFT report, found {len(pdf.pages)} page(s).")

        page1 = pdf.pages[0]
        page2 = pdf.pages[1]

        # --- 1. FULL TEXT METADATA EXTRACTION (Page 1) ---
        full_text_p1 = page1.extract_text() or ""

        lines = [line.strip()
                 for line in full_text_p1.split("\n") if line.strip()]
        for line in lines:
            if "ID: #" in line:
                record["patient_name"] = line.split("ID:")[0].strip()
                dob_match = re.search(r"\((\d{1,2}/\d{1,2}/\d{4})\)", line)
                if dob_match:
                    record["date_of_birth"] = dob_match.group(1)
                break

        dos_match = re.search(
            r"Test Date:\s*(\d{1,2}/\d{1,2}/\d{4})", full_text_p1)
        if dos_match:
            record["date_of_service"] = dos_match.group(1)

        smoker_match = re.search(
            r"Smoker\s*:?\s*(Yes|No|Former)", full_text_p1, re.IGNORECASE)
        if smoker_match:
            record["smoker"] = smoker_match.group(1).capitalize()

        # --- 2. INTERPRETATION DIAGRAM PROCESSING (Page 2) ---
        text_p2 = page2.extract_text() or ""
        detected_box: InterpretationType = "normal"
        box_found = False

        if "System Summary Interpretation" in text_p2:
            summary_text = text_p2.split(
                "System Summary Interpretation")[-1].strip().lower()
            for phrase, token in text_mapping.items():
                if phrase in summary_text:
                    detected_box = token
                    box_found = True
                    break

        if not box_found:
            target_boxes = [
                {"token": "normal", "x0": 300, "x1": 360, "top": 530, "bottom": 570},
                {"token": "pvd", "x0": 350, "x1": 420, "top": 580, "bottom": 630},
                {"token": "cw_neuro_disorder", "x0": 420,
                    "x1": 490, "top": 580, "bottom": 630},
                {"token": "ild_pn", "x0": 490, "x1": 560, "top": 580, "bottom": 630},
                {"token": "asthma_cb", "x0": 560,
                    "x1": 630, "top": 580, "bottom": 630},
                {"token": "emphysema", "x0": 630,
                    "x1": 710, "top": 580, "bottom": 630}
            ]

            shaded_elements = [r for r in page2.rects if r.get(
                "non_stroke_color") is not None]
            for shape in shaded_elements:
                for box in target_boxes:
                    if (box["x0"] <= shape["x0"] <= box["x1"]) and (box["top"] <= shape["top"] <= box["bottom"]):
                        detected_box = box["token"]  # type: ignore
                        box_found = True
                        break

        record["highlighted_box"] = detected_box

    logger.info("PDF parsed, data found: %s",
                json.dumps(record, ensure_ascii=False))

    # --- 3. WRITING TO TEMPLATE ---
    template_file_path = Path(TEMPLATE_NAME)
    if not template_file_path.is_file():
        logger.error("Template file does not exist: %s", TEMPLATE_NAME)
        raise FileNotFoundError(
            f"Template file '{TEMPLATE_NAME}' missing inside container workspace.")

    writer = PdfWriter()
    template_reader = PdfReader(TEMPLATE_NAME)
    raw_results_reader: PdfReader | None = None

    try:
        writer.append(template_reader)

        patient_data = {"patient_name": "",
                        "date_of_service": "", "date_of_birth": ""}
        check_boxes = {
            "results_cw_neuro_disorder": "/Off", "results_ild_pn": "/Off", "results_asthma_cb": "/Off",
            "results_emphysema": "/Off", "refer_pulmo": "/Off", "refer_smoking": "/Off",
            "results_pvd": "/Off", "results_normal": "/Off"
        }

        formatted_dos = datetime.strptime(record.get(
            "date_of_service"), "%m/%d/%Y").strftime("%m/%d/%Y") if record.get("date_of_service") else ""
        formatted_dob = datetime.strptime(record.get(
            "date_of_birth"), "%m/%d/%Y").strftime("%m/%d/%Y") if record.get("date_of_birth") else ""

        patient_data |= {
            "patient_name": record.get("patient_name").title() if record.get("patient_name") else "",
            "date_of_service": formatted_dos,
            "date_of_birth": formatted_dob,
        }

        check_boxes |= {
            "refer_smoking": "/Yes" if record.get("smoker") == "Yes" else "/Off"}

        match record["highlighted_box"]:
            case "normal":
                check_boxes |= {"results_normal": "/Yes"}
            case "pvd":
                check_boxes |= {"results_pvd": "/Yes", "refer_pulmo": "/Yes"}
            case "cw_neuro_disorder":
                check_boxes |= {
                    "results_cw_neuro_disorder": "/Yes", "refer_pulmo": "/Yes"}
            case "ild_pn":
                check_boxes |= {"results_ild_pn": "/Yes", "refer_pulmo": "/Yes"}
            case "asthma_cb":
                check_boxes |= {"results_asthma_cb": "/Yes", "refer_pulmo": "/Yes"}
            case "emphysema":
                check_boxes |= {"results_emphysema": "/Yes", "refer_pulmo": "/Yes"}

        form_data = patient_data | check_boxes

        writer.get_fields()

        writer.update_page_form_field_values(
            writer.pages[0],
            form_data,
            auto_regenerate=False,
            flatten=True
        )

        for page in writer.pages:
            if "/Annots" in page:
                for annot in page["/Annots"]:
                    annot_obj = annot.get_object()
                    if annot_obj.get("/t") in check_boxes:
                        field_name = annot_obj["/T"]
                        target_value = check_boxes[field_name]
                        annot_obj.update({
                            NameObject("/V"): NameObject(target_value),
                            NameObject("/AS"): NameObject(target_value)
                        })

        try:
            writer.set_need_appearances_writer()
        except AttributeError:
            if "/AcroForm" in writer._root_object:
                writer._root_object["/AcroForm"].update({
                    NameObject("/NeedAppearances"): BooleanObject(True)
                })

        # Strip widget annotations
        writer.remove_annotations(subtypes="/Widget")

        # Append raw input pages directly onto the end of our current working context
        input_stream.seek(0)
        raw_results_reader = PdfReader(input_stream)
        writer.append(raw_results_reader)

        # Secure document permissions
        permissions = UserAccessPermissions.PRINT | UserAccessPermissions.EXTRACT
        owner_pw = "heather"

        writer.encrypt(
            user_password="",
            owner_password=owner_pw,
            permissions_flag=permissions
        )

        # Keep small responses in memory, but spill larger PDFs to a temp file.
        output_stream = tempfile.SpooledTemporaryFile(max_size=8 * 1024 * 1024, mode="w+b")
        writer.write(output_stream)
        output_stream.seek(0)
    finally:
        template_reader.close()
        if raw_results_reader is not None:
            raw_results_reader.close()

    logger.info("PDTR successfully prepended to PFT result PDF.")
    return output_stream


@app.route('/process-pdf', methods=['POST'])
def upload_file():
    if 'file' not in request.files:
        logger.info("No file part in request")
        return abort(400, "No file part in the request")

    file = request.files['file']
    if file.filename == '':
        logger.info("No file selector")
        return abort(400, "No file selected")

    if file and file.filename.endswith('.pdf'):
        original_filename = secure_filename(file.filename)
        logger.info("Received PDF: %s", original_filename)

        try:
            # Execute logic sequentially using the upload stream directly.
            processed_file_stream = process_pft_pdf(file.stream)

            logger.info("Returning processed PDF file: %s", original_filename)
            response = send_file(
                processed_file_stream,
                mimetype='application/pdf',
                as_attachment=True,
                download_name=original_filename
            )

            @response.call_on_close
            def close_processed_file_stream() -> None:
                processed_file_stream.close()
                gc.collect()

            return response

        except Exception as e:
            logger.exception("Error during execution processing pipeline")
            return abort(500, f"Error processing file: {str(e)}")
    return abort(400, "Invalid file format. Only PDF files are accepted.")


@app.route('/', methods=['GET'])
def index():
    # Simple, clean HTML page with an upload form for manual testing
    return """
    <!DOCTYPE html>
    <html lang="en">
    <head>
        <meta charset="UTF-8">
        <meta name="viewport" content="width=device-width, initial-scale=1.0">
        <title>PFT Result PDF Processing Service</title>
        <style>
            body { font-family: -apple-system, sans-serif; line-height: 1.6; max-width: 600px; margin: 40px auto; padding: 0 20px; color: #333; }
            h1 { color: #0066cc; }
            .card { background: #f4f4f9; padding: 20px; border-radius: 8px; border: 1px solid #ddd; }
            input[type="file"] { margin: 15px 0; display: block; }
            button { background: #0066cc; color: white; border: none; padding: 10px 15px; border-radius: 4px; cursor: pointer; }
            button:hover { background: #0052a3; }
        </style>
    </head>
    <body>
        <h1>PFT Report Processor</h1>
        <p>This automated API service parses PFT reports and appends formatted data templates.</p>
        <div class="card">
            <h3>Manual Test Upload</h3>
            <form action="/process-pdf" method="post" enctype="multipart/form-data">
                <label for="file">Select PFT Result PDF:</label>
                <input type="file" id="file" name="file" accept=".pdf" required>
                <button type="submit">Upload and Process</button>
            </form>
        </div>
    </body>
    </html>
    """


if __name__ == "__main__":
    app.run(host='0.0.0.0', port=5000)