import asyncio
from llama_index.core.workflow import Context
from llama_index.core.agent.workflow import AgentWorkflow
from tools.tavily_search_tool import (
    search_web, record_notes, write_report, review_report,
    SearchWebArgs, RecordNotesArgs, WriteReportArgs, ReviewReportArgs
)

async def test_tools():
    """Test all tools with the new Pydantic signatures."""
    print("Testing tools with new Pydantic model arguments...")
    
    # Create a simple workflow for context
    workflow = AgentWorkflow(agents=[], root_agent=None)
    ctx = Context(workflow)
    
    # Initialize state
    await ctx.set("state", {
        "research_notes": {},
        "report_content": "Not written yet.",
        "review": "Review required.",
    })
    
    print("\n1. Testing search_web...")
    try:
        search_args = SearchWebArgs(query="artificial intelligence history")
        search_result = await search_web(search_args)
        print(f"✅ search_web worked! Result length: {len(search_result)}")
        print(f"Preview: {search_result[:200]}...")
    except Exception as e:
        print(f"❌ search_web failed: {e}")
    
    print("\n2. Testing record_notes...")
    try:
        notes_args = RecordNotesArgs(
            notes="Test research notes about AI history",
            notes_title="AI History Overview"
        )
        notes_result = await record_notes(ctx, notes_args)
        print(f"✅ record_notes worked! Result: {notes_result}")
        
        # Check state
        state = await ctx.get("state")
        print(f"State after notes: {list(state.keys())}")
        print(f"Research notes: {state.get('research_notes', {})}")
    except Exception as e:
        print(f"❌ record_notes failed: {e}")
    
    print("\n3. Testing write_report...")
    try:
        report_args = WriteReportArgs(
            report_content="""# Artificial Intelligence History

## Introduction
This is a test report about AI history.

## Early Development
AI began in the 1950s with researchers like Alan Turing.

## Modern Era
Today, AI includes machine learning and deep learning.

## Conclusion
AI continues to evolve rapidly.""",
            title="Test AI History Report"
        )
        report_result = await write_report(ctx, report_args)
        print(f"✅ write_report worked! Result: {report_result}")
        
        # Check state
        state = await ctx.get("state")
        print(f"Report content length: {len(state.get('report_content', ''))}")
        print(f"Has structured report: {'structured_report' in state}")
    except Exception as e:
        print(f"❌ write_report failed: {e}")
    
    print("\n4. Testing review_report...")
    try:
        review_args = ReviewReportArgs(review="APPROVED: The report looks good!")
        review_result = await review_report(ctx, review_args)
        print(f"✅ review_report worked! Result: {review_result}")
        
        # Check final state
        state = await ctx.get("state")
        print(f"Final review: {state.get('review', 'No review')}")
    except Exception as e:
        print(f"❌ review_report failed: {e}")
    
    print("\n5. Final state check...")
    final_state = await ctx.get("state")
    print(f"Final state keys: {list(final_state.keys())}")
    print(f"Research notes count: {len(final_state.get('research_notes', {}))}")
    print(f"Report written: {final_state.get('report_content', 'Not written') != 'Not written yet.'}")

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