"""Examples of using StructuredStore and StructuredQueryHelper.

This file demonstrates practical usage patterns for storing and querying
extracted records using the StructuredStore class.

Run this file from the project root:
    cd /home/neil/Documents/DSOL
    python examples/structured_store_usage.py

Or run as a module:
    python -m examples.structured_store_usage
"""

import asyncio
import sys
from pathlib import Path
from uuid import uuid4

# Add project root to Python path
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
from sqlalchemy.orm import sessionmaker

from src.extraction.models import ExtractedRecord
from src.knowledge.structured_store import StructuredStore
from src.knowledge.query_helper import StructuredQueryHelper


async def example_1_basic_storage():
    """Example 1: Basic storage of extracted records."""
    print("\n" + "="*60)
    print("EXAMPLE 1: Basic Storage")
    print("="*60)

    # Get database URL from environment or use default
    import os
    database_url = os.getenv(
        "DATABASE_URL",
        "postgresql+asyncpg://postgres:postgres@localhost:5432/dsol"
    )

    print(f"📡 Connecting to: {database_url.split('@')[1]}")  # Hide password

    # Create async engine
    engine = create_async_engine(database_url, echo=False)

    # Create session
    async_session = sessionmaker(
        engine, class_=AsyncSession, expire_on_commit=False
    )

    async with async_session() as session:
        # IMPORTANT: Create a document first (foreign key requirement)
        from src.db.models import Document

        document = Document(
            id=uuid4(),
            filename="example.xlsx",
            file_path="/tmp/example.xlsx",
            status="processing"
        )
        session.add(document)
        await session.commit()
        await session.refresh(document)

        document_id = document.id
        print(f"📄 Created document: {document_id}")

        records = [
            ExtractedRecord(
                content={
                    "Item Code": "2024-4_0001",
                    "Description": "Test Item 1",
                    "Quantity": "100",
                    "Screen": "○"
                },
                headers=["Item Code", "Description", "Quantity", "Screen"],
                _source={
                    "document_id": str(document_id),
                    "filename": "example.xlsx",
                    "sheet": "Sheet1",
                    "table_id": 1,
                    "row": 10,
                    "col_range": "A:D"
                }
            ),
            ExtractedRecord(
                content={
                    "Item Code": "2024-4_0002",
                    "Description": "Test Item 2",
                    "Quantity": "200",
                    "Screen": "×"
                },
                headers=["Item Code", "Description", "Quantity", "Screen"],
                _source={
                    "document_id": str(document_id),
                    "filename": "example.xlsx",
                    "sheet": "Sheet1",
                    "table_id": 1,
                    "row": 11,
                    "col_range": "A:D"
                }
            ),
        ]

        # Store records
        store = StructuredStore(session)
        count = await store.store_records(records, document_id)

        print(f"\n✅ Stored {count} records")
        print(f"   Document ID: {document_id}")

    await engine.dispose()


async def example_2_query_exact_match():
    """Example 2: Query records by exact field match."""
    print("\n" + "="*60)
    print("EXAMPLE 2: Query by Exact Field Match")
    print("="*60)

    engine = create_async_engine(
        "postgresql+asyncpg://dsol_user:dsol_password@localhost:5432/dsol",
        echo=False
    )

    async_session = sessionmaker(
        engine, class_=AsyncSession, expire_on_commit=False
    )

    async with async_session() as session:
        # Assume document_id from Example 1
        document_id = uuid4()  # Replace with actual document_id

        # Query by exact Item Code
        query_helper = StructuredQueryHelper(session)
        records = await query_helper.find_by_exact_field(
            document_id,
            "Item Code",
            "2024-4_0001"
        )

        print(f"\n🔍 Found {len(records)} records with Item Code = '2024-4_0001'")
        for record in records:
            print(f"   Content: {record.content}")

    await engine.dispose()


async def example_3_query_containment():
    """Example 3: Query records using JSONB containment."""
    print("\n" + "="*60)
    print("EXAMPLE 3: Query by JSONB Containment")
    print("="*60)

    engine = create_async_engine(
        "postgresql+asyncpg://dsol_user:dsol_password@localhost:5432/dsol",
        echo=False
    )

    async_session = sessionmaker(
        engine, class_=AsyncSession, expire_on_commit=False
    )

    async with async_session() as session:
        document_id = uuid4()  # Replace with actual document_id

        # Find all records where Screen = "○"
        query_helper = StructuredQueryHelper(session)
        records = await query_helper.find_by_containment(
            document_id,
            {"Screen": "○"}
        )

        print(f"\n🔍 Found {len(records)} records with Screen = '○'")
        for record in records:
            print(f"   Item: {record.content.get('Item Code')}")
            print(f"   Screen: {record.content.get('Screen')}")

    await engine.dispose()


async def example_4_query_by_sheet():
    """Example 4: Query all records from a specific sheet."""
    print("\n" + "="*60)
    print("EXAMPLE 4: Query by Sheet")
    print("="*60)

    engine = create_async_engine(
        "postgresql+asyncpg://dsol_user:dsol_password@localhost:5432/dsol",
        echo=False
    )

    async_session = sessionmaker(
        engine, class_=AsyncSession, expire_on_commit=False
    )

    async with async_session() as session:
        document_id = uuid4()  # Replace with actual document_id

        # Get all records from Sheet1
        query_helper = StructuredQueryHelper(session)
        records = await query_helper.find_by_sheet(
            document_id,
            "Sheet1"
        )

        print(f"\n📊 Found {len(records)} records from Sheet1")
        print(f"   Document ID: {document_id}")

    await engine.dispose()


async def example_5_batch_large_dataset():
    """Example 5: Store large dataset (automatic batching)."""
    print("\n" + "="*60)
    print("EXAMPLE 5: Batch Storage (Large Dataset)")
    print("="*60)

    engine = create_async_engine(
        "postgresql+asyncpg://dsol_user:dsol_password@localhost:5432/dsol",
        echo=False
    )

    async_session = sessionmaker(
        engine, class_=AsyncSession, expire_on_commit=False
    )

    async with async_session() as session:
        document_id = uuid4()

        # Create 1200 records (will be 3 batches of 500, 500, 200)
        records = []
        for i in range(1200):
            record = ExtractedRecord(
                content={
                    "Item Code": f"2024-4_{i:04d}",
                    "Description": f"Batch Item {i}",
                    "Quantity": str(100 + i)
                },
                headers=["Item Code", "Description", "Quantity"],
                _source={
                    "document_id": str(document_id),
                    "filename": "large_batch.xlsx",
                    "sheet": "Data",
                    "table_id": 1,
                    "row": 10 + i,
                    "col_range": "A:C"
                }
            )
            records.append(record)

        # Store all records (automatically batched)
        store = StructuredStore(session)
        count = await store.store_records(records, document_id)

        print(f"\n✅ Stored {count} records in batches of 500")
        print(f"   Total batches: {(count + 499) // 500}")
        print(f"   Check logs for performance metrics!")

    await engine.dispose()


async def example_6_idempotent_retry():
    """Example 6: Demonstrate idempotent retry (ON CONFLICT)."""
    print("\n" + "="*60)
    print("EXAMPLE 6: Idempotent Retry (ON CONFLICT)")
    print("="*60)

    engine = create_async_engine(
        "postgresql+asyncpg://dsol_user:dsol_password@localhost:5432/dsol",
        echo=False
    )

    async_session = sessionmaker(
        engine, class_=AsyncSession, expire_on_commit=False
    )

    async with async_session() as session:
        document_id = uuid4()

        # Create a record
        record = ExtractedRecord(
            content={
                "Item Code": "2024-4_RETRY",
                "Description": "Original Description",
                "Version": "1"
            },
            headers=["Item Code", "Description", "Version"],
            _source={
                "document_id": str(document_id),
                "filename": "retry.xlsx",
                "sheet": "Sheet1",
                "table_id": 1,
                "row": 5,
                "col_range": "A:C"
            }
        )

        # First insert
        store = StructuredStore(session)
        count1 = await store.store_records([record], document_id)
        print(f"\n✅ First insert: {count1} record")

        # Modify and retry (same document_id, sheet, row)
        record.content["Description"] = "Updated Description"
        record.content["Version"] = "2"

        count2 = await store.store_records([record], document_id)
        print(f"✅ Second insert (update): {count2} record")
        print(f"   ON CONFLICT handled - record updated, not duplicated")

    await engine.dispose()


async def main():
    """Run all examples."""
    print("\n" + "🚀"*30)
    print("StructuredStore Usage Examples")
    print("🚀"*30)

    print("\n⚠️  Prerequisites:")
    print("  1. PostgreSQL running (docker-compose up -d postgres)")
    print("  2. Database migrations applied (alembic upgrade head)")
    print("  3. DATABASE_URL environment variable set")
    print()

    try:
        # Run examples sequentially
        await example_1_basic_storage()
        # Uncomment these after first example succeeds:
        # await example_2_query_exact_match()
        # await example_3_query_containment()
        # await example_4_query_by_sheet()
        # await example_5_batch_large_dataset()
        # await example_6_idempotent_retry()

        print("\n" + "✅"*30)
        print("All examples completed!")
        print("✅"*30 + "\n")

    except Exception as e:
        print("\n" + "❌"*30)
        print(f"Error running examples: {e}")
        print("❌"*30)
        print("\n💡 Troubleshooting:")
        print("  1. Check if PostgreSQL is running:")
        print("     docker ps | grep postgres")
        print("  2. Check database credentials:")
        print("     echo $DATABASE_URL")
        print("  3. Run migrations:")
        print("     alembic upgrade head")
        print()
        raise


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