import os
import logging

import fitz  # PyMuPDF

from app.config import settings

logger = logging.getLogger(__name__)


def get_page_count(file_path: str) -> int:
    doc = fitz.open(file_path)
    count = len(doc)
    doc.close()
    return count


def extract_text_by_page(file_path: str) -> dict[int, str]:
    """Extract text from every page. Returns {page_num: text} (1-indexed)."""
    doc = fitz.open(file_path)
    pages = {}
    for i in range(len(doc)):
        text = doc[i].get_text()
        if text.strip():
            pages[i + 1] = text
    doc.close()
    return pages


def extract_text_for_range(file_path: str, start: int, end: int) -> str:
    """Extract text for a page range (1-indexed, inclusive)."""
    doc = fitz.open(file_path)
    texts = []
    for i in range(start - 1, min(end, len(doc))):
        text = doc[i].get_text()
        if text.strip():
            texts.append(f"--- Page {i + 1} ---\n{text}")
    doc.close()
    return "\n\n".join(texts)


def extract_images_from_page(file_path: str, page_num: int, document_id: int) -> list[str]:
    """Extract images from a specific page. Returns list of saved image paths."""
    image_dir = os.path.join(settings.UPLOAD_DIR, "images", f"doc_{document_id}")
    os.makedirs(image_dir, exist_ok=True)

    saved_paths = []
    try:
        doc = fitz.open(file_path)
        if page_num < 1 or page_num > len(doc):
            doc.close()
            return []

        page = doc[page_num - 1]
        images = page.get_images(full=True)

        for img_idx, img_info in enumerate(images):
            xref = img_info[0]
            try:
                base_image = doc.extract_image(xref)
                if not base_image:
                    continue

                image_ext = base_image.get("ext", "png")
                image_bytes = base_image["image"]

                # Skip tiny images (likely icons/bullets)
                if len(image_bytes) < 1000:
                    continue

                filename = f"page_{page_num}_img_{img_idx}.{image_ext}"
                filepath = os.path.join(image_dir, filename)

                with open(filepath, "wb") as f:
                    f.write(image_bytes)

                # Return relative path for serving
                rel_path = f"images/doc_{document_id}/{filename}"
                saved_paths.append(rel_path)
            except Exception as e:
                logger.warning(f"Failed to extract image {img_idx} from page {page_num}: {e}")
                continue

        doc.close()
    except Exception as e:
        logger.warning(f"Image extraction failed for page {page_num}: {e}")

    return saved_paths


def extract_all_images(file_path: str, document_id: int, start_page: int = 1, end_page: int | None = None) -> dict[int, list[str]]:
    """Extract images from a page range. Returns {page_num: [image_paths]}."""
    doc = fitz.open(file_path)
    if end_page is None:
        end_page = len(doc)
    doc.close()

    result = {}
    for page_num in range(start_page, end_page + 1):
        images = extract_images_from_page(file_path, page_num, document_id)
        if images:
            result[page_num] = images
    return result