#!/usr/bin/env python3
"""
Script to ingest PDF and PowerPoint files from CustomerDocument directory through the API.
"""

import requests
import time
from pathlib import Path

API_BASE = "http://localhost:8000"
API_KEY = "962aed8f4137f5fc20d718f488e0ab15ce49a11572f76cad614ccb4a453d1702"

# Base directory for customer documents
CUSTOMER_DOC_DIR = Path(__file__).parent / "CustomerDocument"


def find_pdf_files() -> list[Path]:
    """Find all PDF files (.pdf) in CustomerDocument directory."""
    pdf_files = []
    pdf_files.extend(CUSTOMER_DOC_DIR.rglob("*.pdf"))
    return sorted(pdf_files)


def find_pptx_files() -> list[Path]:
    """Find all PowerPoint files (.ppt, .pptx, .pptm) in CustomerDocument directory."""
    pptx_files = []
    for ext in ["*.ppt", "*.pptx", "*.pptm"]:
        pptx_files.extend(CUSTOMER_DOC_DIR.rglob(ext))
    return sorted(pptx_files)


def check_health():
    """Check if API is healthy."""
    try:
        response = requests.get(f"{API_BASE}/health", timeout=5)
        health = response.json()
        print(f"API Health: {health['status']}")
        print(f"  Database: {health['dependencies']['database']['status']}")
        print(f"  Redis: {health['dependencies']['redis']['status']}")
        print(f"  Milvus: {health['dependencies']['milvus']['status']}")
        return health['status'] == 'healthy' or health['dependencies']['milvus']['status'] == 'healthy'
    except Exception as e:
        print(f"Health check failed: {e}")
        return False


def upload_file(file_path: str):
    """Upload a file to the API."""
    path = Path(file_path)

    if not path.exists():
        print(f"❌ File not found: {file_path}")
        return None

    print(f"\n📤 Uploading: {path.name}")
    print(f"   Size: {path.stat().st_size / 1024:.2f} KB")

    try:
        with open(file_path, 'rb') as f:
            files = {'file': (path.name, f)}
            headers = {'X-API-Key': API_KEY}
            response = requests.post(
                f"{API_BASE}/documents",
                files=files,
                headers=headers,
                timeout=30
            )

        if response.status_code in [200, 202]:
            result = response.json()
            doc_id = result.get('document_id')
            print(f"✅ Uploaded successfully")
            print(f"   Document ID: {doc_id}")
            return doc_id
        else:
            print(f"❌ Upload failed: {response.status_code}")
            print(f"   Error: {response.text}")
            return None

    except Exception as e:
        print(f"❌ Upload error: {e}")
        return None


def check_document_status(doc_id: str, timeout: int = 120):
    """Check document processing status."""
    print(f"⏳ Checking status for {doc_id}...")

    start_time = time.time()
    while time.time() - start_time < timeout:
        try:
            response = requests.get(f"{API_BASE}/documents/{doc_id}", timeout=5)
            if response.status_code == 200:
                doc = response.json()
                status = doc.get('status', 'unknown')
                print(f"   Status: {status}", end='\r')

                if status == 'completed':
                    print(f"\n✅ Processing complete!")
                    return True
                elif status == 'failed':
                    print(f"\n❌ Processing failed!")
                    return False

            time.sleep(2)
        except Exception as e:
            print(f"   Error checking status: {e}")
            time.sleep(2)

    print(f"\n⏱️  Timeout waiting for processing")
    return False


def main():
    print("=" * 70)
    print("DSOL - PDF and PowerPoint File Ingestion")
    print("=" * 70)

    # Check API health
    print("\n1. Checking API health...")
    if not check_health():
        print("API not fully healthy, but continuing...")

    # Find all PDF and PowerPoint files
    pdf_files = find_pdf_files()
    pptx_files = find_pptx_files()
    print(f"\n2. Found files to ingest:")
    print(f"   PDF files: {len(pdf_files)}")
    print(f"   PowerPoint files: {len(pptx_files)}")
    print(f"   Total: {len(pdf_files) + len(pptx_files)}")

    # Upload PDF files
    print(f"\n3. Uploading PDF files...")
    pdf_docs = []
    for file_path in pdf_files[:5]:  # Limit to first 5 PDFs for testing
        doc_id = upload_file(str(file_path))
        if doc_id:
            pdf_docs.append(doc_id)
        time.sleep(0.5)  # Small delay between uploads

    # Upload PowerPoint files
    print(f"\n4. Uploading PowerPoint files...")
    pptx_docs = []
    for file_path in pptx_files[:5]:  # Limit to first 5 PowerPoints for testing
        doc_id = upload_file(str(file_path))
        if doc_id:
            pptx_docs.append(doc_id)
        time.sleep(0.5)  # Small delay between uploads

    # Wait for processing
    print(f"\n5. Waiting for processing to complete...")
    print(f"   PDF documents uploaded: {len(pdf_docs)}")
    print(f"   PowerPoint documents uploaded: {len(pptx_docs)}")
    print(f"   Total documents uploaded: {len(pdf_docs) + len(pptx_docs)}")

    time.sleep(5)  # Give Celery workers time to pick up tasks

    # Check status of PDF files
    print(f"\n6. Checking PDF files status:")
    pdf_success = 0
    for doc_id in pdf_docs:
        if check_document_status(doc_id, timeout=300):  # Longer timeout for PDFs
            pdf_success += 1

    # Check status of PowerPoint files
    print(f"\n7. Checking PowerPoint files status:")
    pptx_success = 0
    for doc_id in pptx_docs:
        if check_document_status(doc_id, timeout=120):
            pptx_success += 1

    # Summary
    print("\n" + "=" * 70)
    print("INGESTION SUMMARY")
    print("=" * 70)
    print(f"PDF files found: {len(pdf_files)}")
    print(f"PDF files uploaded: {len(pdf_docs)}")
    print(f"PDF files processed successfully: {pdf_success}/{len(pdf_docs)}")
    print()
    print(f"PowerPoint files found: {len(pptx_files)}")
    print(f"PowerPoint files uploaded: {len(pptx_docs)}")
    print(f"PowerPoint files processed successfully: {pptx_success}/{len(pptx_docs)}")
    print()
    print(f"Total success rate: {(pdf_success + pptx_success)}/{(len(pdf_docs) + len(pptx_docs))}")
    print("=" * 70)


if __name__ == "__main__":
    main()