main.py

import os
import sys
import tempfile
import logging

from utils.file_helpers import allowed_file, get_file_extension
from utils.processors import PDFProcessor, PPTXProcessor

pdf_processor = PDFProcessor()
pptx_processor = PPTXProcessor()

logger = logging.getLogger(__name__)

OUTPUT_FOLDER = "output"
os.makedirs(OUTPUT_FOLDER, exist_ok=True)


def allowed_file(filename):
    return filename.lower().endswith((".pdf", ".pptx"))


def get_file_extension(filename):
    return filename.rsplit(".", 1)[-1].lower() if "." in filename else ""


def _process_pdf_local(upload_path, filename):
    output_filename = f"processed_{filename}"
    output_path = os.path.join(OUTPUT_FOLDER, output_filename)

    result = pdf_processor.process(upload_path, output_path, filename)

    if not result["success"]:
        raise RuntimeError(result["error"])

    return {
        "message": result["message"],
        "output_path": output_path if result["has_watermark"] else None,
        "file_type": "pdf",
        "has_watermark": result["has_watermark"],
    }


def _process_pptx_local(upload_path, filename):
    output_filename = f"processed_{filename}"
    output_path = os.path.join(OUTPUT_FOLDER, output_filename)

    result = pptx_processor.process(upload_path, output_path, filename)

    if not result["success"]:
        raise RuntimeError(result["error"])

    return {
        "message": result["message"],
        "output_path": output_path if result["has_watermark"] else None,
        "file_type": "pptx",
        "has_watermark": result["has_watermark"],
    }


def remove_watermark_local(file_path):
    if not os.path.exists(file_path):
        raise ValueError("File does not exist")

    filename = os.path.basename(file_path)

    if not allowed_file(filename):
        raise ValueError("Invalid file type. Only PDF or PPTX allowed.")

    original_extension = get_file_extension(filename)

    if not original_extension:
        raise ValueError("File must have a valid extension (.pdf or .pptx)")

    file_extension = original_extension.lower()

    logger.info(f"Processing file: {filename} (type: {file_extension})")

    with tempfile.NamedTemporaryFile(delete=False, suffix=f".{file_extension}") as temp_input:
        upload_path = temp_input.name

        try:
            with open(file_path, "rb") as f:
                temp_input.write(f.read())
            temp_input.flush()

            if file_extension == "pdf":
                return _process_pdf_local(upload_path, filename)
            elif file_extension == "pptx":
                return _process_pptx_local(upload_path, filename)
            else:
                raise ValueError(f"Unsupported file type: {file_extension}")

        finally:
            try:
                os.unlink(upload_path)
            except Exception:
                pass


def main():
    if len(sys.argv) != 2:
        print(f"Usage: {sys.argv[0]} <file.pdf|file.pptx>")
        sys.exit(1)

    file_path = sys.argv[1]

    try:
        result = remove_watermark_local(file_path)
        print(result["message"])

        if result["output_path"]:
            print(result["output_path"])
        else:
            print("No watermark found.")
            print(file_path)

        sys.exit(0)

    except Exception as e:
        print(f"Error: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()