import io
from pathlib import Path

from pdf2image import convert_from_bytes
from PIL import Image, ImageChops, ImageDraw, ImageFont

ALLOWED_IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png"}
ALLOWED_PDF_EXTENSION = ".pdf"
ALLOWED_SPREADSHEET_EXTENSIONS = {".xlsx", ".xls", ".csv"}

# DPI for PDF rendering — 200 gives good quality for dense invoices
_PDF_DPI = 200

# Max pages to process per PDF — covers most multi-page invoices
_PDF_MAX_PAGES = 6


def get_extension(filename: str) -> str:
    return Path(filename or "").suffix.lower()


def is_pdf(filename: str) -> bool:
    return get_extension(filename) == ALLOWED_PDF_EXTENSION


def is_image(filename: str) -> bool:
    return get_extension(filename) in ALLOWED_IMAGE_EXTENSIONS


def is_spreadsheet(filename: str) -> bool:
    return get_extension(filename) in ALLOWED_SPREADSHEET_EXTENSIONS


def validate_file(filename: str) -> None:
    if not (is_pdf(filename) or is_image(filename) or is_spreadsheet(filename)):
        raise ValueError(
            "Only PDF, JPG, JPEG, PNG, XLSX, XLS, and CSV files are supported"
        )


def to_png_bytes(content: bytes, filename: str) -> bytes:
    """
    Convert any supported file type to a single PNG image for the LLM.
    For multi-page PDFs, all pages are stitched vertically into one tall image
    so the model sees the complete document in a single pass.
    """
    if is_pdf(filename):
        return _pdf_to_stitched_png(content)

    if is_spreadsheet(filename):
        return _spreadsheet_to_png(content, filename)

    pil_image = Image.open(io.BytesIO(content))
    return _pil_to_png(_normalize_document_image(pil_image))


def _pdf_to_stitched_png(content: bytes) -> bytes:
    """
    Render all pages of a PDF (up to _PDF_MAX_PAGES) at _PDF_DPI and
    stitch them vertically into a single PNG. This lets the LLM see the
    full document — including totals on the last page — in one image.
    For strategic stock sheets (very dense), use higher DPI.
    """
    # Try higher DPI first for dense documents, fall back if too large
    dpi = _PDF_DPI
    pages = convert_from_bytes(
        content,
        first_page=1,
        last_page=_PDF_MAX_PAGES,
        dpi=dpi,
    )
    if not pages:
        raise ValueError("PDF has no readable pages")

    if len(pages) == 1:
        return _pil_to_png(_normalize_document_image(pages[0]))

    # Convert all pages to RGB PIL images
    pil_pages = []
    for page in pages:
        if page.mode in ("RGBA", "P"):
            page = page.convert("RGB")
        pil_pages.append(page)

    # Stitch vertically — use the widest page as the canvas width
    max_width = max(p.width for p in pil_pages)
    total_height = sum(p.height for p in pil_pages) + (len(pil_pages) - 1) * 8  # 8px gap

    stitched = Image.new("RGB", (max_width, total_height), color="#e8ecf4")
    y_offset = 0
    for page in pil_pages:
        # Centre narrower pages
        x_offset = (max_width - page.width) // 2
        stitched.paste(page, (x_offset, y_offset))
        y_offset += page.height + 8

    return _pil_to_png(_normalize_document_image(stitched))


def _spreadsheet_to_png(content: bytes, filename: str) -> bytes:
    """
    Render a spreadsheet as a PNG image so the LLM can read it.
    Uses openpyxl for xlsx/xls and the csv module for csv files.
    Renders the first 80 rows × 25 columns as a plain table image.
    """
    ext = get_extension(filename)
    rows: list[list[str]] = []

    if ext in {".xlsx", ".xls"}:
        try:
            import openpyxl  # type: ignore
            wb = openpyxl.load_workbook(io.BytesIO(content), read_only=True, data_only=True)
            ws = wb.active
            for i, row in enumerate(ws.iter_rows(values_only=True)):  # type: ignore
                if i >= 80:
                    break
                rows.append([str(cell) if cell is not None else "" for cell in row[:25]])
        except Exception:
            rows = _parse_csv_bytes(content)
    elif ext == ".csv":
        rows = _parse_csv_bytes(content)

    if not rows:
        raise ValueError("Spreadsheet appears to be empty")

    return _rows_to_png(rows)


def _parse_csv_bytes(content: bytes) -> list[list[str]]:
    import csv

    text = content.decode("utf-8", errors="replace")
    reader = csv.reader(io.StringIO(text))
    rows = []
    for i, row in enumerate(reader):
        if i >= 80:
            break
        rows.append([cell[:40] for cell in row[:25]])
    return rows


def _rows_to_png(rows: list[list[str]]) -> bytes:
    """Render a list of rows as a monospaced table PNG."""
    col_width = 150
    row_height = 22
    padding = 10
    font_size = 11

    num_cols = max(len(r) for r in rows) if rows else 1
    width = num_cols * col_width + padding * 2
    height = len(rows) * row_height + padding * 2

    img = Image.new("RGB", (width, height), color="#ffffff")
    draw = ImageDraw.Draw(img)

    try:
        font = ImageFont.truetype("/System/Library/Fonts/Helvetica.ttc", font_size)
        header_font = ImageFont.truetype("/System/Library/Fonts/Helvetica.ttc", font_size)
    except Exception:
        try:
            font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", font_size)
            header_font = font
        except Exception:
            font = ImageFont.load_default()
            header_font = font

    for row_idx, row in enumerate(rows):
        y = padding + row_idx * row_height
        bg = "#f4f7fd" if row_idx == 0 else ("#f9fbff" if row_idx % 2 == 0 else "#ffffff")
        draw.rectangle([0, y, width, y + row_height], fill=bg)

        for col_idx, cell in enumerate(row):
            x = padding + col_idx * col_width
            text = str(cell)[:20]
            f = header_font if row_idx == 0 else font
            draw.text((x + 4, y + 4), text, fill="#1f2740", font=f)

        draw.line([0, y + row_height, width, y + row_height], fill="#dfe5f0", width=1)

    for col_idx in range(num_cols + 1):
        x = padding + col_idx * col_width
        draw.line([x, 0, x, height], fill="#dfe5f0", width=1)

    return _pil_to_png(img)


def _pil_to_png(image: Image.Image) -> bytes:
    if image.mode in ("RGBA", "P"):
        image = image.convert("RGB")
    buffer = io.BytesIO()
    image.save(buffer, format="PNG")
    return buffer.getvalue()


def _normalize_document_image(image: Image.Image) -> Image.Image:
    """
    Improve document readability for vision extraction:
    - trim large white page margins
    - auto-rotate portrait pages whose non-white content is clearly landscape
    """
    if image.mode in ("RGBA", "P"):
        image = image.convert("RGB")

    image = _trim_white_margins(image)

    bbox = _content_bbox(image)
    if bbox:
        content_w = bbox[2] - bbox[0]
        content_h = bbox[3] - bbox[1]
        # Some scanned sheets are embedded sideways inside a portrait PDF page.
        # Rotate when the page is portrait but the actual content block is wide.
        if image.height > image.width and content_w > int(content_h * 1.1):
            rotated = image.rotate(270, expand=True)
            image = _trim_white_margins(rotated)

    return image


def _trim_white_margins(image: Image.Image, margin: int = 12) -> Image.Image:
    bbox = _content_bbox(image)
    if not bbox:
        return image

    left = max(bbox[0] - margin, 0)
    top = max(bbox[1] - margin, 0)
    right = min(bbox[2] + margin, image.width)
    bottom = min(bbox[3] + margin, image.height)
    return image.crop((left, top, right, bottom))


def _content_bbox(image: Image.Image):
    bg = Image.new(image.mode, image.size, "white")
    diff = ImageChops.difference(image, bg).convert("L")
    # Ignore tiny scan speckles while keeping table lines and text.
    mask = diff.point(lambda p: 255 if p > 18 else 0)
    return mask.getbbox()