"""Test the agent-based simulation."""

import asyncio
import os

# Add src to path
import sys
sys.path.insert(0, "src")

from mirofish_simulator import StudentAgent, run_simulation

# Test question
QUESTION = {
    "text": "What is the primary function of the Electoral College?",
    "options": [
        "To directly elect members of Congress",
        "To formally elect the President and Vice President",
        "To approve Supreme Court nominations",
        "To ratify constitutional amendments"
    ],
    "correct_answer": "B",
    "subject": "AP Government",
}


async def test_single_agent():
    """Test a single student agent."""
    print("=" * 60)
    print("TEST 1: Single Agent (Grade 5 ESL Student)")
    print("=" * 60)

    agent = StudentAgent(grade=5, archetype="esl_student")
    response = await agent.answer(QUESTION, correct_answer=QUESTION["correct_answer"])

    print(f"\nQuestion: {QUESTION['text'][:50]}...")
    print(f"Correct Answer: {QUESTION['correct_answer']}")
    print()
    print(f"Agent Selected: {response.selected}")
    print(f"Is Correct: {response.is_correct}")
    print(f"Confidence: {response.confidence}")
    print()
    print(f"Reasoning: {response.reasoning[:200]}...")
    print()
    print(f"Knowledge Used: {response.knowledge_used}")
    print(f"Knowledge Gaps: {response.knowledge_gaps}")
    print(f"Words Not Understood: {response.words_not_understood}")
    print()
    print(f"What they thought question asked:")
    print(f"  {response.what_they_thought_question_asked[:100]}...")

    return response


async def test_multiple_archetypes():
    """Test different archetypes on same question."""
    print("\n" + "=" * 60)
    print("TEST 2: Multiple Archetypes (Grade 5)")
    print("=" * 60)

    archetypes = [
        "honors_overachiever",
        "quiet_thinker",
        "esl_student",
        "class_clown",
        "disengaged_but_smart",
    ]

    results = []
    for arch in archetypes:
        agent = StudentAgent(grade=5, archetype=arch)
        response = await agent.answer(QUESTION, correct_answer=QUESTION["correct_answer"])
        results.append((arch, response))
        print(f"\n{arch}:")
        print(f"  Selected: {response.selected} ({'✓' if response.is_correct else '✗'})")
        print(f"  Confidence: {response.confidence}")
        print(f"  Reasoning: {response.reasoning[:80]}...")

    correct = sum(1 for _, r in results if r.is_correct)
    print(f"\nOverall: {correct}/{len(results)} correct")

    return results


async def test_grade_comparison():
    """Test same archetype at different grades."""
    print("\n" + "=" * 60)
    print("TEST 3: Grade Comparison (Quiet Thinker)")
    print("=" * 60)

    grades = [5, 8, 11]

    for grade in grades:
        agent = StudentAgent(grade=grade, archetype="quiet_thinker")
        response = await agent.answer(QUESTION, correct_answer=QUESTION["correct_answer"])
        print(f"\nGrade {grade}:")
        print(f"  Selected: {response.selected} ({'✓' if response.is_correct else '✗'})")
        print(f"  Confidence: {response.confidence}")
        print(f"  Knowledge gaps: {response.knowledge_gaps[:2]}")


async def test_simulation():
    """Test full simulation with population."""
    print("\n" + "=" * 60)
    print("TEST 4: Full Simulation (Grade 8, 15 students)")
    print("=" * 60)

    result = await run_simulation(
        question=QUESTION,
        grade=8,
        population_size=15,
    )

    print(f"\nOverall Accuracy: {result.accuracy:.0%}")
    print(f"Total Responses: {result.total_responses}")

    print(f"\nResponse Distribution:")
    for option, count in sorted(result.response_distribution.items()):
        marker = " <-- correct" if option == result.correct_answer else ""
        print(f"  {option}: {count} students{marker}")

    print(f"\nBy Archetype:")
    for arch, data in sorted(result.by_archetype.items(), key=lambda x: -x[1].accuracy):
        print(f"  {arch}: {data.accuracy:.0%} ({data.correct}/{data.count})")

    print(f"\nCommon Knowledge Gaps:")
    for gap in result.common_knowledge_gaps[:5]:
        print(f"  - {gap}")

    return result


async def main():
    """Run all tests."""
    print("\n🧪 MiroFish Agent-Based Simulation Tests\n")

    # Check for API key
    if not os.environ.get("OPENAI_API_KEY"):
        print("⚠️  No OPENAI_API_KEY found - tests will use fallback (random) mode")
        print("   Set OPENAI_API_KEY for real agent behavior\n")

    await test_single_agent()
    await test_multiple_archetypes()
    await test_grade_comparison()
    await test_simulation()

    print("\n" + "=" * 60)
    print("✅ All tests completed!")
    print("=" * 60)


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