"""
PDF Scan -> Google Drive OCR -> Text Extraction
==============================================
Upload scanned PDFs to Google Drive (auto-converts to Google Docs with OCR),
export the recognized text, save as .txt files, and clean up.

Usage:
    python ocr_drive.py                    # Process all PDFs in current directory
    python ocr_drive.py file1.pdf file2.pdf # Process specific files
"""

import os
import sys
import io
import json
import time
import glob
import mimetypes
from pathlib import Path

from google.oauth2 import service_account
from googleapiclient.discovery import build
from googleapiclient.http import MediaFileUpload, MediaIoBaseDownload

# --- CONFIG ---
SCRIPT_DIR = Path(__file__).parent
SERVICE_ACCOUNT_FILE = SCRIPT_DIR / "ocr_service_account.json"
SCOPES = ["https://www.googleapis.com/auth/drive.file"]
WAIT_SECONDS = 5  # Wait for Google to process OCR
GOOGLE_DOC_MIME = "application/vnd.google-apps.document"
DOCX_EXPORT_MIME = "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
TEXT_EXPORT_MIME = "text/plain"


def get_drive_service():
    """Authenticate and return Google Drive service."""
    if not SERVICE_ACCOUNT_FILE.exists():
        print(f"[ERROR] Khong tim thay file credentials: {SERVICE_ACCOUNT_FILE}")
        sys.exit(1)

    creds = service_account.Credentials.from_service_account_file(
        str(SERVICE_ACCOUNT_FILE), scopes=SCOPES
    )
    return build("drive", "v3", credentials=creds)


def _guess_mime_type(file_path: str, source_mime_type: str = None) -> str:
    """Best-effort MIME detection for upload to Google Drive."""
    if source_mime_type:
        source_mime_type = source_mime_type.strip().lower()
        if source_mime_type and source_mime_type != "application/octet-stream":
            return source_mime_type

    guessed, _ = mimetypes.guess_type(file_path)
    if guessed:
        return guessed

    return "application/octet-stream"


def ocr_file(
    drive,
    file_path: str,
    source_mime_type: str = None,
    filename: str = None,
    export_mode: str = "text",
):
    """
    Upload a file to Google Drive -> convert to Google Docs (OCR/export) -> export data -> cleanup.

    export_mode:
      - "text": return UTF-8 string
      - "docx": return DOCX bytes
    """
    local_name = os.path.basename(file_path)
    display_name = filename or local_name
    upload_mime = _guess_mime_type(file_path, source_mime_type)

    print(f"\n{'='*60}")
    print(f"[FILE] Processing: {display_name}")
    print(f"   Size: {os.path.getsize(file_path) / 1024:.0f} KB")
    print(f"   MIME: {upload_mime}")

    doc_id = None
    try:
        # Step 1: Upload source file -> Google Docs conversion
        print(f"   [UPLOAD] Uploading to Google Drive...")
        file_metadata = {"name": display_name, "mimeType": GOOGLE_DOC_MIME}
        media = MediaFileUpload(file_path, mimetype=upload_mime)
        uploaded = (
            drive.files()
            .create(body=file_metadata, media_body=media, fields="id")
            .execute()
        )
        doc_id = uploaded.get("id")
        print(f"   [OK] Uploaded! Doc ID: {doc_id}")

        # Step 2: Wait for Google OCR processing
        print(f"   [WAIT] Waiting {WAIT_SECONDS}s for OCR processing...")
        time.sleep(WAIT_SECONDS)

        # Step 3: Export from Google Docs
        if export_mode == "docx":
            export_mime = DOCX_EXPORT_MIME
            print(f"   [EXPORT] Exporting DOCX...")
        else:
            export_mime = TEXT_EXPORT_MIME
            print(f"   [EXPORT] Exporting text...")

        request = drive.files().export_media(fileId=doc_id, mimeType=export_mime)
        fh = io.BytesIO()
        downloader = MediaIoBaseDownload(fh, request)
        done = False
        while not done:
            status, done = downloader.next_chunk()
            if status:
                print(f"   [DOWNLOAD] {int(status.progress() * 100)}%")

        payload = fh.getvalue()
        if export_mode == "docx":
            print(f"   [OK] Exported {len(payload):,} bytes DOCX")
            return payload

        text_content = payload.decode("utf-8").strip()
        print(f"   [OK] Extracted {len(text_content):,} characters")
        return text_content

    except Exception as e:
        print(f"   [ERROR] {e}")
        if export_mode == "docx":
            return b""
        return ""

    finally:
        # Step 4: Cleanup - delete from Google Drive
        if doc_id:
            try:
                drive.files().delete(fileId=doc_id).execute()
                print(f"   [CLEANUP] Deleted Drive file")
            except Exception as e:
                print(f"   [WARN] Cleanup failed: {e}")


def ocr_pdf(drive, pdf_path: str) -> str:
    """Backward-compatible wrapper for legacy callers."""
    return ocr_file(drive, pdf_path, source_mime_type="application/pdf")


def save_text(pdf_path: str, text: str) -> str:
    """Save extracted text to a .txt file with the same name."""
    txt_path = os.path.splitext(pdf_path)[0] + ".txt"
    with open(txt_path, "w", encoding="utf-8") as f:
        f.write(text)
    print(f"   [SAVED] {os.path.basename(txt_path)}")
    return txt_path


def main():
    # Determine which files to process
    if len(sys.argv) > 1:
        pdf_files = [f for f in sys.argv[1:] if f.lower().endswith(".pdf")]
    else:
        pdf_files = glob.glob(str(SCRIPT_DIR / "*.pdf"))

    if not pdf_files:
        print("[ERROR] Khong tim thay file PDF nao!")
        print(f"   Thu muc: {SCRIPT_DIR}")
        sys.exit(1)

    print(f"[START] Google Drive OCR - {len(pdf_files)} file(s)")
    print(f"   Credentials: {SERVICE_ACCOUNT_FILE}")

    # Authenticate
    drive = get_drive_service()
    print("[OK] Ket noi Google Drive thanh cong!")

    # Process each PDF
    results = []
    for pdf_path in pdf_files:
        text = ocr_file(drive, pdf_path, source_mime_type="application/pdf")
        if text:
            txt_path = save_text(pdf_path, text)
            results.append((pdf_path, txt_path, len(text)))

            # Preview first 500 chars
            print(f"\n   [PREVIEW] 500 ky tu dau:")
            print(f"   {'-'*50}")
            preview = text[:500].replace("\n", "\n   ")
            print(f"   {preview}")
            print(f"   {'-'*50}")
        else:
            print(f"   [FAIL] Khong trich xuat duoc text tu {os.path.basename(pdf_path)}")

    # Summary
    print(f"\n{'='*60}")
    print(f"[SUMMARY] KET QUA TONG HOP")
    print(f"{'='*60}")
    print(f"   Tong file: {len(pdf_files)}")
    print(f"   Thanh cong: {len(results)}")
    print(f"   That bai: {len(pdf_files) - len(results)}")
    for pdf_path, txt_path, chars in results:
        print(f"   [OK] {os.path.basename(pdf_path)} -> {chars:,} ky tu")
    print()


if __name__ == "__main__":
    main()