import gradio as gr
try:
    import fitz  # from PyMuPDF
except ImportError:
    raise ImportError("PyMuPDF is not installed. Make sure to include 'pymupdf' in your requirements.txt.")

import numpy as np
import pandas as pd
from PIL import Image
import io

def extract_pdf_images(pdf_bytes):
    doc = fitz.open(stream=pdf_bytes, filetype="pdf")
    images = []
    for page in doc:
        pix = page.get_pixmap(matrix=fitz.Matrix(2, 2))
        img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
        images.append(img)
    return images

def calculate_area(drawing_data, page_number, pdf_file, area_type):
    if drawing_data is None or "layers" not in drawing_data:
        return None, "No drawing found.", None

    try:
        coords = []
        for layer in drawing_data["layers"]:
            for shape in layer["paths"]:
                if shape["type"] == "polygon":
                    coords.extend([[pt["x"], pt["y"]] for pt in shape["points"]])

        coords = np.array(coords)
        if len(coords) < 3:
            return None, "Draw a closed polygon to calculate area.", None

        area_px = 0.5 * np.abs(np.dot(coords[:, 0], np.roll(coords[:, 1], 1)) -
                               np.dot(coords[:, 1], np.roll(coords[:, 0], 1)))
        area_sft = area_px / (96**2) * 144

        rd_label = f"Page {page_number + 1}"
        result_df = pd.DataFrame([{
            "RD": rd_label,
            "Area Type": area_type,
            "Area (sq.ft)": round(area_sft, 2)
        }])

        excel_io = io.BytesIO()
        with pd.ExcelWriter(excel_io, engine="openpyxl") as writer:
            result_df.to_excel(writer, index=False, sheet_name="Area Summary")
        excel_bytes = excel_io.getvalue()

        return f"Estimated {area_type}: {area_sft:.2f} sq.ft", None, excel_bytes
    except Exception as e:
        return None, str(e), None

with gr.Blocks() as demo:
    gr.Markdown("## 📐 Canal Cross Section Area Calculator (Gradio Version)")
    gr.Markdown("Upload a multi-page PDF of canal cross sections. Select a page, draw polygon on the image, and compute area (in sqft).")

    pdf_input = gr.File(label="📄 Upload PDF", type="binary")
    area_type = gr.Radio(choices=["Cutting Area", "Filling Area"], value="Cutting Area", label="Area Type")
    page_slider = gr.Slider(minimum=0, maximum=0, step=1, label="Select Page")
    image_display = gr.Image(label="Cross Section Page", interactive=False)
    drawing = gr.Sketchpad(label="Draw Area on Image", shape="polygon")
    calculate_btn = gr.Button("🧮 Calculate Area")
    result_text = gr.Textbox(label="Result")
    download_excel = gr.File(label="📥 Download Excel", interactive=False)
    error_output = gr.Textbox(label="Errors", visible=False)

    def load_pdf(pdf_file):
        if pdf_file is None:
            return gr.Image.update(value=None), gr.update(maximum=0), None
        images = extract_pdf_images(pdf_file)
        return images[0], gr.update(maximum=len(images) - 1), images

    def update_page(images, page_index):
        return images[page_index]

    pdf_file_state = gr.State()
    images_state = gr.State()

    pdf_input.change(load_pdf, inputs=[pdf_input], outputs=[image_display, page_slider, images_state])
    page_slider.change(update_page, inputs=[images_state, page_slider], outputs=image_display)

    calculate_btn.click(
        calculate_area,
        inputs=[drawing, page_slider, pdf_input, area_type],
        outputs=[result_text, error_output, download_excel]
    )

demo.launch()