#!/usr/bin/env python3
"""
Quick verification script to test the new fund relationship schema
"""

import sys

sys.path.insert(0, "/home/oluwasanmi/Documents/Work/MKD/anton_wireframe/preprocessor")

from models import FundTable, InvestmentStageTable, SectorTable, get_db_session


def test_fund_relationships():
    """Test the new fund relationship schema"""
    db = get_db_session()

    print("🧪 Testing Fund Relationship Schema\n")

    # Test 1: Check investment stages
    print("1️⃣ Investment Stages:")
    stages = db.query(InvestmentStageTable).all()
    print(f"   Found {len(stages)} stages:")
    for stage in stages[:5]:
        print(f"   - {stage.name}")
    print()

    # Test 2: Check fund with relationships
    print("2️⃣ Sample Fund with Relationships:")
    fund = db.query(FundTable).filter(FundTable.fund_name.isnot(None)).first()

    if fund:
        print(f"   Fund: {fund.fund_name}")
        print(f"   Geographic Focus: {fund.geographic_focus}")

        print(f"   Investment Stages ({len(fund.investment_stages)}):")
        for stage in fund.investment_stages[:3]:
            print(f"   - {stage.name}")

        print(f"   Sectors ({len(fund.sectors)}):")
        for sector in fund.sectors[:3]:
            print(f"   - {sector.name}")
    else:
        print("   No funds found")
    print()

    # Test 3: Check association tables
    print("3️⃣ Association Table Stats:")

    # Count fund-stage relationships
    from sqlalchemy import text

    result = db.execute(text("SELECT COUNT(*) FROM fund_investment_stages"))
    stage_count = result.scalar()
    print(f"   Fund-Stage relationships: {stage_count}")

    # Count fund-sector relationships
    result = db.execute(text("SELECT COUNT(*) FROM fund_sectors"))
    sector_count = result.scalar()
    print(f"   Fund-Sector relationships: {sector_count}")
    print()

    # Test 4: Query funds by stage
    print("4️⃣ Query Test - Funds with 'Series A' stage:")
    series_a_funds = (
        db.query(FundTable)
        .join(FundTable.investment_stages)
        .filter(InvestmentStageTable.name.ilike("%Series A%"))
        .limit(3)
        .all()
    )

    print(f"   Found {len(series_a_funds)} funds:")
    for fund in series_a_funds:
        print(f"   - {fund.fund_name or 'Unnamed'}")
        stages = [s.name for s in fund.investment_stages]
        print(f"     Stages: {', '.join(stages)}")
    print()

    # Test 5: Query funds by sector
    print("5️⃣ Query Test - Funds investing in first sector:")
    first_sector = db.query(SectorTable).first()
    if first_sector:
        sector_funds = (
            db.query(FundTable)
            .join(FundTable.sectors)
            .filter(SectorTable.id == first_sector.id)
            .limit(3)
            .all()
        )

        print(f"   Sector: {first_sector.name}")
        print(f"   Found {len(sector_funds)} funds:")
        for fund in sector_funds:
            print(f"   - {fund.fund_name or 'Unnamed'}")
    print()

    # Test 6: Geographic focus string search
    print("6️⃣ Query Test - Funds with Europe in geographic focus:")
    europe_funds = (
        db.query(FundTable)
        .filter(FundTable.geographic_focus.ilike("%Europe%"))
        .limit(3)
        .all()
    )

    print(f"   Found {len(europe_funds)} funds:")
    for fund in europe_funds:
        print(f"   - {fund.fund_name or 'Unnamed'}")
        print(f"     Geographic Focus: {fund.geographic_focus}")
    print()

    print("✅ All tests completed successfully!")
    db.close()


if __name__ == "__main__":
    try:
        test_fund_relationships()
    except Exception as e:
        print(f"❌ Error: {e}")
        import traceback

        traceback.print_exc()