from __future__ import annotations

import argparse
import logging
from pathlib import Path

import fitz

logging.basicConfig(level=logging.INFO, format="%(message)s")
cwd = Path.cwd()
logger = logging.getLogger(__name__)


def parse_page_ranges(range_str: str, total_pages: int) -> list[int]:
    """Parse page range string and return list of page numbers (1-indexed)."""
    pages = []
    for part in range_str.split(","):
        part = part.strip()
        if not part:
            continue
        if "-" in part:
            start, end = part.split("-")
            start = int(start) if start else 1
            end = int(end) if end else total_pages
            pages.extend(range(start, end + 1))
        else:
            pages.append(int(part))
    return pages


def split_by_number(input_file: Path, output_file: Path, number: int):
    """Split PDF into specified number of parts evenly."""
    doc = fitz.open(input_file)
    total_pages = doc.page_count
    base_pages = total_pages // number
    remainder = total_pages % number

    logger.debug(
        f"Total pages: {total_pages}, Splitting into {number} parts, {base_pages} base pages per part, {remainder} extra pages"
    )

    current_page = 0
    for i in range(number):
        # First 'remainder' parts get one extra page
        pages_in_this_part = base_pages + (1 if i < remainder else 0)

        if current_page >= total_pages:
            logger.debug(f"Skipping part {i + 1}: no more pages remaining")
            continue

        end_page = min(current_page + pages_in_this_part, total_pages)

        part_file = output_file.parent / f"{output_file.stem}_part{i + 1}{output_file.suffix}"
        part_doc = fitz.open()

        for page_num in range(current_page, end_page):
            part_doc.insert_pdf(doc, from_page=page_num, to_page=page_num)

        part_doc.save(part_file)
        part_doc.close()
        logger.info(f"Created part {i + 1}: {part_file} (pages {current_page + 1}-{end_page})")

        current_page = end_page

    doc.close()


def split_by_size(input_file: Path, output_file: Path, size: int):
    """Split PDF into parts with specified page size."""
    doc = fitz.open(input_file)
    total_pages = doc.page_count

    logger.debug(f"Total pages: {total_pages}, Splitting with {size} pages per part")

    part = 0
    start_page = 0

    while start_page < total_pages:
        end_page = min(start_page + size, total_pages)
        part_file = output_file.parent / f"{output_file.stem}_part{part + 1}{output_file.suffix}"
        part_doc = fitz.open()

        for page_num in range(start_page, end_page):
            part_doc.insert_pdf(doc, from_page=page_num, to_page=page_num)

        part_doc.save(part_file)
        part_doc.close()
        logger.info(f"Created part {part + 1}: {part_file} (pages {start_page + 1}-{end_page})")

        start_page = end_page
        part += 1

    doc.close()


def split_by_range(input_file: Path, output_file: Path, range_str: str):
    """Extract specific pages from PDF based on range string."""
    doc = fitz.open(input_file)
    total_pages = doc.page_count

    pages = parse_page_ranges(range_str, total_pages)
    pages = [p - 1 for p in pages if 1 <= p <= total_pages]  # Convert to 0-indexed

    if not pages:
        logger.error("No valid pages found in the specified range")
        doc.close()
        return

    # Remove duplicates while preserving order
    pages = sorted(set(pages))

    logger.debug(f"Extracting pages: {[p + 1 for p in pages]}")

    new_doc = fitz.open()
    for page_num in pages:
        new_doc.insert_pdf(doc, from_page=page_num, to_page=page_num)

    new_doc.save(output_file)
    new_doc.close()
    doc.close()
    logger.info(f"Created output file: {output_file} ({len(pages)} pages)")


def main():
    parser = argparse.ArgumentParser(description="Split PDF files")
    parser.add_argument("input", help="Input PDF file")
    parser.add_argument("output", nargs="?", help="Output PDF file (optional for -n and -s modes)")
    parser.add_argument("-o", "--output-dir", default=str(cwd), help="Output directory (default: current directory)")
    parser.add_argument("-f", "--output-format", help="Output file format pattern, e.g., 'split_{part:02d}.pdf'")
    parser.add_argument("-v", "--verbose", action="store_true", help="Verbose output")

    # Split by number, size, or range
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("-n", "--number", type=int, help="Number of splits")
    group.add_argument("-s", "--size", type=int, default=1, help="Size of each split in pages")
    group.add_argument("-r", "--range", type=str, help="Range of pages to extract, e.g., '1,2,4-10,15-20,25-'")

    args = parser.parse_args()

    if args.verbose:
        logger.setLevel(logging.DEBUG)

    output_dir = Path(args.output_dir)
    if not output_dir.is_dir():
        logger.error(f"Output directory {args.output_dir} does not exist, please check the path.")
        return

    input_file = Path(args.input)
    if not input_file.is_file():
        logger.error(f"Input file {args.input} does not exist, please check the path.")
        return

    # For -n and -s modes, output is optional and defaults to base name with suffix
    # For -r mode, output is required
    if args.range and not args.output:
        logger.error("Output file is required for -r/--range mode")
        return

    if not args.range:
        output_file = output_dir / (input_file.stem + "_split.pdf") if not args.output else Path(args.output)
    else:
        output_file = Path(args.output)

    logger.info(f"Start splitting {input_file}")
    if args.number:
        split_by_number(input_file, output_file, args.number)
    elif args.size:
        split_by_size(input_file, output_file, args.size)
    elif args.range:
        split_by_range(input_file, output_file, args.range)
    else:
        logger.error("Please specify either -n, -s, or -r")
        return
