#!/usr/bin/env python3
"""Test script to verify ORM changes work correctly."""

import os
import sys

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from src.local_deep_research.utilities.db_utils import get_db_session
from src.local_deep_research.web.database.models import (
    Base,
    Research,
    ResearchMode,
    ResearchStatus,
)
from src.local_deep_research.web.services.research_service import (
    get_research_strategy,
    save_research_strategy,
)


def test_orm_changes():
    """Test that the ORM changes work correctly."""
    print("Testing ORM changes...")

    # Get the session and create tables in the actual database being used
    session = get_db_session()

    # Get the engine from the session
    engine = session.bind

    # Create all tables if they don't exist
    Base.metadata.create_all(engine)
    print("✓ Tables created successfully")

    # Test save_research_strategy and get_research_strategy
    # First create a test research
    try:
        test_research = Research(
            query="Test query",
            status=ResearchStatus.COMPLETED,
            mode=ResearchMode.QUICK,
        )
        session.add(test_research)
        session.commit()
        research_id = test_research.id
        print(f"✓ Created test research with ID: {research_id}")
    finally:
        session.close()

    # Test saving a strategy
    test_strategy = "source-based"
    save_research_strategy(research_id, test_strategy)
    print(f"✓ Saved strategy '{test_strategy}' for research {research_id}")

    # Test retrieving the strategy
    retrieved_strategy = get_research_strategy(research_id)
    assert retrieved_strategy == test_strategy, (
        f"Expected '{test_strategy}', got '{retrieved_strategy}'"
    )
    print(f"✓ Retrieved strategy correctly: '{retrieved_strategy}'")

    # Test updating a strategy
    new_strategy = "iterative"
    save_research_strategy(research_id, new_strategy)
    print(f"✓ Updated strategy to '{new_strategy}'")

    retrieved_strategy = get_research_strategy(research_id)
    assert retrieved_strategy == new_strategy, (
        f"Expected '{new_strategy}', got '{retrieved_strategy}'"
    )
    print(f"✓ Retrieved updated strategy correctly: '{retrieved_strategy}'")

    # Cleanup - delete the test research record
    session = get_db_session()
    try:
        test_research = session.get(Research, research_id)
        if test_research:
            # Also delete associated research_strategies
            from src.local_deep_research.web.database.models import (
                ResearchStrategy,
            )

            strategies = (
                session.query(ResearchStrategy)
                .filter_by(research_id=research_id)
                .all()
            )
            for strategy in strategies:
                session.delete(strategy)
            session.delete(test_research)
            session.commit()
            print("✓ Cleaned up test data")
    finally:
        session.close()

    print("\n✅ All ORM tests passed!")


if __name__ == "__main__":
    test_orm_changes()
