import io
import tempfile
from typing import List

import pypdfium2

from surya.debug.draw import draw_polys_on_image, draw_bboxes_on_image

from PIL import Image
from surya.settings import settings
from vllm_tools import predictors_vllm, ocr_vllm


predictors = predictors_vllm


def rescale_bbox(bbox, source_size, target_size):
    width_ratio = target_size[0] / source_size[0]
    height_ratio = target_size[1] / source_size[1]
    return [
        bbox[0] * width_ratio,
        bbox[1] * height_ratio,
        bbox[2] * width_ratio,
        bbox[3] * height_ratio,
    ]


def expand_bbox(bbox, margin=5):
    return [
        max(0, int(bbox[0]) - margin),
        max(0, int(bbox[1]) - margin),
        int(bbox[2]) + margin,
        int(bbox[3]) + margin,
    ]



def page_counter(pdf_file):
    doc = open_pdf(pdf_file)
    doc_len = len(doc)
    doc.close()
    return doc_len

def ocr_errors(pdf_file, page_count, sample_len=512, max_samples=10, max_pages=15):
    from pdftext.extraction import plain_text_output

    with tempfile.NamedTemporaryFile(suffix=".pdf") as f:
        f.write(pdf_file.getvalue())
        f.seek(0)

        # Sample the text from the middle of the PDF
        page_middle = page_count // 2
        page_range = range(
            max(page_middle - max_pages, 0), min(page_middle + max_pages, page_count)
        )
        text = plain_text_output(f.name, page_range=page_range)

    sample_gap = len(text) // max_samples
    if len(text) == 0 or sample_gap == 0:
        return "This PDF has no text or very little text", ["no text"]

    if sample_gap < sample_len:
        sample_gap = sample_len

    # Split the text into samples for the model
    samples = []
    for i in range(0, len(text), sample_gap):
        samples.append(text[i : i + sample_len])

    results = predictors["ocr_error"](samples)
    label = "This PDF has good text."
    if results.labels.count("bad") / len(results.labels) > 0.2:
        label = "This PDF may have garbled or bad OCR text."
    return label, results.labels


def text_detection(img):
    text_pred = predictors["detection"]([img])[0]
    text_polygons = [p.polygon for p in text_pred.bboxes]
    det_img = draw_polys_on_image(text_polygons, img.copy())
    return det_img, text_pred


def layout_detection(img):
    pred = predictors["layout"]([img])[0]
    polygons = [p.polygon for p in pred.bboxes]
    labels = [
        f"{p.label}-{p.position}-{round(p.top_k[p.label], 2)}" for p in pred.bboxes
    ]
    layout_img = draw_polys_on_image(
        polygons, img.copy(), labels=labels, label_font_size=18
    )
    return layout_img, pred


def table_recognition(
    img, highres_img, skip_table_detection: bool
):
    if skip_table_detection:
        layout_tables = [(0, 0, highres_img.size[0], highres_img.size[1])]
        table_imgs = [highres_img]
    else:
        _, layout_pred = layout_detection(img)
        layout_tables_lowres = [
            line.bbox
            for line in layout_pred.bboxes
            if line.label in ["Table", "TableOfContents"]
        ]
        table_imgs = []
        layout_tables = []
        for tb in layout_tables_lowres:
            highres_bbox = rescale_bbox(tb, img.size, highres_img.size)
            # Slightly expand the box
            highres_bbox = expand_bbox(highres_bbox)
            table_imgs.append(highres_img.crop(highres_bbox))
            layout_tables.append(highres_bbox)

    table_preds = predictors["table_rec"](table_imgs)
    table_img = img.copy()

    for results, table_bbox in zip(table_preds, layout_tables):
        adjusted_bboxes = []
        labels = []
        colors = []

        for item in results.cells:
            adjusted_bboxes.append(
                [
                    (item.bbox[0] + table_bbox[0]),
                    (item.bbox[1] + table_bbox[1]),
                    (item.bbox[2] + table_bbox[0]),
                    (item.bbox[3] + table_bbox[1]),
                ]
            )
            labels.append(item.label)
            if "Row" in item.label:
                colors.append("blue")
            else:
                colors.append("red")
        table_img = draw_bboxes_on_image(
            adjusted_bboxes,
            highres_img,
            labels=labels,
            label_font_size=18,
            color=colors,
        )
    return table_img, table_preds


# Function for OCR
def ocr(
    img: Image.Image,
    highres_img: Image.Image,
    skip_text_detection: bool = False,
    recognize_math: bool = True,
    with_bboxes: bool = True,
):
    return ocr_vllm(
        img,
        highres_img,
        skip_text_detection=skip_text_detection,
        recognize_math=recognize_math,
        with_bboxes=with_bboxes,
    )


def open_pdf(pdf_file):
    stream = io.BytesIO(pdf_file.getvalue())
    return pypdfium2.PdfDocument(stream)


def get_page_image(pdf_file, page_num, dpi=settings.IMAGE_DPI):
    doc = open_pdf(pdf_file)
    renderer = doc.render(
        pypdfium2.PdfBitmap.to_pil,
        page_indices=[page_num - 1],
        scale=dpi / 72,
    )
    png = list(renderer)[0]
    png_image = png.convert("RGB")
    doc.close()
    return png_image


def page_counter(pdf_file):
    doc = open_pdf(pdf_file)
    doc_len = len(doc)
    doc.close()
    return doc_len


import pandas as pd

def bbox_intersection(box1, box2):
    x1 = max(box1[0], box2[0])
    y1 = max(box1[1], box2[1])
    x2 = min(box1[2], box2[2])
    y2 = min(box1[3], box2[3])
    
    if x1 < x2 and y1 < y2:
        return (x1, y1, x2, y2)
    else:
        return None

def area_of_bbox(box):
    return (box[2] - box[0]) * (box[3] - box[1])

def is_bbox_inside(box, parent_box):
    interaction_box = bbox_intersection(box, parent_box)
    if interaction_box is None:
        return False
    return area_of_bbox(interaction_box) / area_of_bbox(box) > 0.5

def center_of_bbox(box):
    return ((box[0] + box[2]) / 2, (box[1] + box[3]) / 2)


def extract_text_from_image(image):
    layout_predictor = predictors["layout"]
    recognition_predictor = predictors["recognition"]

    layout_prediction = layout_predictor([image])[0]
    prediction = recognition_predictor([image], [layout_prediction], full_page=False)[0]

    items = [
        {
            "text": block.html,
            "position": block.reading_order,
            "order_value": block.bbox[1],
        }
        for block in prediction.blocks
        if not block.skipped and not block.error and block.html
    ]
    if not items:
        return ""
    
    df = pd.DataFrame(items)
    df = df.sort_values(by=['position', 'order_value'])
    ds = df.groupby('position').apply(lambda x: "  ".join(x['text'].tolist()))
    full_text = "\n\n".join(ds.to_list())
    
    return full_text
