import sys
import gc
import re
import json
import io
import logging
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 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_bytes: bytes) -> io.BytesIO:
    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"
    }

    # Use a detached context for pdfplumber to read the file metadata
    with pdfplumber.open(io.BytesIO(input_bytes)) 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()

        if full_text_p1:
            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()
        detected_box: InterpretationType = "normal"
        box_found = False

        if text_p2 and "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)
    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
    raw_results_reader = PdfReader(io.BytesIO(input_bytes))
    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
    )

    # Output final joined stream directly to Flask
    output_stream = io.BytesIO()
    writer.write(output_stream)
    output_stream.seek(0)

    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:
            # Read input file directly into memory bytes
            raw_input_bytes = file.read()

            # Execute logic sequentially
            processed_file_stream = process_pft_pdf(raw_input_bytes)

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

        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>
    """


@app.teardown_request
def spi_clear_memory(exception=None):
    """Runs automatically after every single request."""
    # Force collection of unreferenced objects
    gc.collect()


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