#!/usr/bin/env python3
"""
Test script for Enhanced Parser Agent
Tests multimodal extraction: text + figures + tables + embeddings
"""
import asyncio
import sys
from pathlib import Path

# Add parent directory to path
sys.path.insert(0, str(Path(__file__).parent.parent))

from agents.enhanced_parser_agent import EnhancedParserAgent
from agents.base_agent import AgentContext
from services.neo4j_service import Neo4jService
from services.mongo_service import MongoService
from services.llm_service import LLMService
from services.vector_store_service import VectorStoreService
from services.lancedb_service import LanceDBService
from memory.graph_memory import GraphMemory
from utils.helpers import generate_id
from utils.logger import get_logger

logger = get_logger("test_parser")


async def test_enhanced_parser(pdf_path: str) -> dict:
    """Test the enhanced parser on a PDF file."""

    if not Path(pdf_path).exists():
        logger.error(f"PDF file not found: {pdf_path}")
        return {"error": f"File not found: {pdf_path}"}

    logger.info(f"Testing enhanced parser on: {pdf_path}")

    # Initialize services
    try:
        neo4j = Neo4jService()
        mongo = MongoService()
        llm = LLMService()
        embedder = VectorStoreService()
        vector_db = LanceDBService()

        logger.info("Services initialized")
    except Exception as e:
        logger.error(f"Failed to initialize services: {e}")
        return {"error": str(e)}

    # Create context
    job_id = "test_" + generate_id("job")[:8]
    doc_id = generate_id("doc")

    graph_memory = GraphMemory(neo4j=neo4j, embedder=embedder)
    ctx = AgentContext(
        job_id=job_id,
        doc_id=doc_id,
        pdf_path=pdf_path,
        neo4j=neo4j,
        mongo=mongo,
        llm=llm,
        embedder=embedder,
        graph_memory=graph_memory,
    )
    ctx.vector_db = vector_db

    # Run enhanced parser
    parser = EnhancedParserAgent()
    try:
        result = await parser.execute(ctx)
        logger.info(f"Parser result: {result}")
        return result
    except Exception as e:
        logger.error(f"Parser execution failed: {e}")
        return {"error": str(e)}


async def main():
    """Main test function."""

    # Check for PDF argument
    if len(sys.argv) < 2:
        logger.info("Usage: python test_docling_parser.py <pdf_path>")
        logger.info("Example: python test_docling_parser.py ./samples/paper.pdf")
        return

    pdf_path = sys.argv[1]

    logger.info("=" * 60)
    logger.info("ENHANCED PARSER AGENT TEST")
    logger.info("=" * 60)

    result = await test_enhanced_parser(pdf_path)

    logger.info("=" * 60)
    logger.info("TEST RESULTS:")
    logger.info("=" * 60)

    if "error" in result:
        logger.error(f"TEST FAILED: {result['error']}")
    else:
        logger.info(f"✓ Parser: {result.get('parser', 'unknown')}")
        logger.info(f"✓ Title: {result.get('title', 'N/A')}")
        logger.info(f"✓ Sections: {result.get('section_count', 0)}")
        logger.info(f"✓ Figures: {result.get('figure_count', 0)}")
        logger.info(f"✓ Tables: {result.get('table_count', 0)}")
        logger.info(f"✓ Vector indexed: {result.get('vector_indexed', False)}")
        logger.info("\nTEST PASSED ✓")


if __name__ == "__main__":
    asyncio.run(main())