main.py

import os
import pdfplumber
import ollama
import pywhatkit as kit
import time
from docx import Document
from docx.shared import Pt, RGBColor
from docx.enum.text import WD_ALIGN_PARAGRAPH


# UPDATE THESE — REQUIRED

PDF_PATH = "input.pdf"          # 👈 Path to your PDF file
OUTPUT_PATH = "output.docx"     # 👈 Where to save the Word document

# WHATSAPP SETTINGS
WHATSAPP_NUMBER = "+911234567890"  # 👈 Your WhatsApp number (with country code)
SEND_TO_WHATSAPP = True            # 👈 Set to False if you only want the Word doc
WHATSAPP_WAIT_TIME = 15            # 👈 Seconds to wait for WhatsApp Web to load
CLOSE_TAB_AFTER = True             # 👈 Close browser tab after sending?


#  UPDATE THESE — OLLAMA SETTINGS

MODEL_NAME = "llama3.2"         # 👈 Run `ollama list` in terminal to confirm name
                                 #    Common names: llama3.2, llama3, mistral, phi3

PROMPT = (                       # 👈 Change this to control what Ollama does
    "You are a document analyst. Analyze the following PDF text and provide:"
    "1. SUMMARY: A brief 2-3 sentence overview of what this document is about"
    "2. KEY ACTION POINTS: A numbered list of the most important actions, tasks, or takeaways from the document"
    "3. IMPORTANT DETAILS: Any critical dates, names, numbers, or deadlines mentioned"
    "Keep it concise and clear. Only include what is actually present in the document."
)

# ============================================================
# ✏️  UPDATE THESE — DOCX FORMATTING (optional)
# ============================================================

# Use the PDF filename as the document title
PDF_FILENAME = os.path.basename(PDF_PATH)
DOC_TITLE = os.path.splitext(PDF_FILENAME)[0]  # Extracts name without .pdf extension
FONT_NAME = "Calibri"              # 👈 Font for body text (Arial, Times New Roman, etc.)
FONT_SIZE = 11                     # 👈 Font size for body text (in points)


#NO CHANGES NEEDED BELOW THIS LINE

def extract_text_page_by_page(pdf_path):
    """
    Reads every page of the PDF individually.
    Returns a list of texts, one for each page.
    """
    if not os.path.exists(pdf_path):
        raise FileNotFoundError(f"PDF not found at: {pdf_path}")

    pages_content = []
    with pdfplumber.open(pdf_path) as pdf:
        total_pages = len(pdf.pages)
        print(f"   Found {total_pages} page(s) in PDF")

        for i, page in enumerate(pdf.pages):
            page_text = page.extract_text()
            if page_text and page_text.strip():
                pages_content.append({
                    "number": i + 1,
                    "text": page_text.strip()
                })
            else:
                print(f"   ⚠️  Page {i + 1} has no extractable text (might be a scanned image)")

    return pages_content


def process_with_ollama(text, prompt):
    """
    Sends the extracted PDF text to your local Ollama model.
    Uses the official `ollama` Python library.
    Ollama must be running on your machine (ollama serve).
    """
    print(f"   Model  : {MODEL_NAME}")
    print(f"   Prompt : {prompt[:70]}...")

    response = ollama.chat(
        model=MODEL_NAME,
        messages=[
            {
                "role": "user",
                "content": f"{prompt}\n\n{text}"
            }
        ]
    )
    return response.message.content


def send_to_whatsapp(phone, message):
    """
    Sends a message via WhatsApp Web using pywhatkit.
    Requires you to be logged into WhatsApp Web on your default browser.
    """
    # Basic format check (must start with + and contain digits)
    if not phone.startswith("+") or not phone[1:].isdigit():
        raise ValueError(f"Number {phone} is not in correct format (must start with + and include country code)")

    print(f"   📤 Sending message to {phone}...")
    kit.sendwhatmsg_instantly(
        phone_no=phone,
        message=message,
        wait_time=WHATSAPP_WAIT_TIME,
        tab_close=CLOSE_TAB_AFTER,
        close_time=3
    )
    print("   ✅ Message sent!")


def write_to_docx(content, output_path):
    """
    Writes the Ollama-processed text into a formatted .docx Word document.
    - Lines starting with '--- Page N ---' become Word headings
    - Everything else becomes body paragraphs
    """
    doc = Document()

    # --- Document title ---
    title = doc.add_heading(DOC_TITLE, level=0)
    title.alignment = WD_ALIGN_PARAGRAPH.CENTER

    # --- Add a thin separator line after title ---
    doc.add_paragraph("_" * 60)

    # --- Write each line as a paragraph or heading ---
    for line in content.split("\n"):
        line = line.strip()

        if not line:
            continue  # skip blank lines

        elif line.startswith("---") and line.endswith("---"):
            # Convert "--- Page 1 ---" markers into Word Heading 1
            heading_text = line.strip("- ").strip()
            doc.add_heading(heading_text, level=1)

        elif line.startswith("#"):
            # Support markdown-style headings from Ollama output
            level = min(line.count("#", 0, 6), 4)
            heading_text = line.lstrip("#").strip()
            doc.add_heading(heading_text, level=level)

        else:
            # Regular body paragraph
            para = doc.add_paragraph()
            run = para.add_run(line)
            run.font.name = FONT_NAME
            run.font.size = Pt(FONT_SIZE)

    doc.save(output_path)


# ============================================================
# MAIN PIPELINE
# ============================================================

def main():
    print("\n" + "=" * 55)
    print("  DocuMate AI Pipeline")
    print("=" * 55)

    # Step 1 — Extract text from PDF individually
    print("\n📄 Step 1: Extracting text page-by-page...")
    try:
        pages = extract_text_page_by_page(PDF_PATH)
    except FileNotFoundError as e:
        print(f"❌ Error: {e}")
        return

    if not pages:
        print("❌ No text was extracted from the PDF.")
        return

    print(f"   ✅ Extracted text from {len(pages)} page(s)")

    combined_summary = ""

    # Step 2 & 3 — Process each page and send to WhatsApp
    for page in pages:
        print(f"\n🤖 Processing Page {page['number']}...")
        
        try:
            # Get summary for THIS page
            summary = process_with_ollama(page['text'], PROMPT)
            
            page_summary_header = f"📄 PAGE {page['number']} SUMMARY:\n"
            full_message = f"{page_summary_header}{summary}"
            
            combined_summary += f"--- Page {page['number']} ---\n{summary}\n\n"

            # Send to WhatsApp
            if SEND_TO_WHATSAPP:
                try:
                    send_to_whatsapp(WHATSAPP_NUMBER, full_message)
                    # Small delay between pages to ensure the previous tab closes/completes
                    time.sleep(2)
                except Exception as e:
                    print(f"\n❌ Error: Number {WHATSAPP_NUMBER} does not exist or could not be reached.")
                    print("   Check the number and try again. Extraction will now stop.")
                    return  # Stop the entire script
            
        except Exception as e:
            print(f"❌ Error during Page {page['number']} processing: {e}")
            return  # Stop processing if any serious error occurs

    # Step 4 — Write everything back to .docx
    print("\n📝 Step 4: Finalizing Word document...")
    write_to_docx(combined_summary, OUTPUT_PATH)
    print(f"   ✅ Full summary saved to: {os.path.abspath(OUTPUT_PATH)}")

    print("\n✅ All pages processed!\n")


if __name__ == "__main__":
    main()