"""Test if agents actually USE the constrained knowledge.

The problem: LLM still knows the right answer even when told it doesn't.

Solution test: Give the agent WRONG information and see if they use it.
If they ignore the wrong info and pick the right answer, the constraint is broken.
If they pick the wrong answer based on wrong info, the constraint works.
"""

import asyncio
import json
import os
import sys
sys.path.insert(0, "src")

from openai import AsyncOpenAI


QUESTION = {
    "text": "How many electoral votes are needed to win the presidency?",
    "options": ["218", "270", "300", "435"],
    "correct_answer": "B",  # 270 is correct
}


async def test_with_wrong_knowledge():
    """Test if agent uses explicitly wrong knowledge."""

    client = AsyncOpenAI()

    # Give the agent WRONG information
    system_prompt = """You are a grade 5 student taking a test.

=== WHAT YOU KNOW (these are FACTS you learned in school) ===
- The President is elected by electoral votes
- You need 300 electoral votes to win (you learned this in class)
- There are 538 total electoral votes

=== IMPORTANT ===
- You MUST answer based on what you learned above
- These facts are what YOU know - use them
- Do NOT use any other knowledge

Answer the question based ONLY on what you learned above."""

    user_prompt = f"""Question: {QUESTION['text']}

Options:
A) 218
B) 270
C) 300
D) 435

Based on what you learned in class (300 electoral votes needed), which answer is correct?

Return JSON: {{"selected": "A/B/C/D", "reasoning": "why you picked this"}}"""

    print("=" * 60)
    print("TEST: Does the agent USE wrong knowledge?")
    print("=" * 60)
    print(f"\nWrong fact given: 'You need 300 electoral votes to win'")
    print(f"Correct answer: B (270)")
    print(f"Expected if constraint works: C (300)")
    print()

    response = await client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ],
        temperature=0.7,
    )

    result_text = response.choices[0].message.content.strip()
    print(f"Raw response:\n{result_text}\n")

    # Parse
    if "```" in result_text:
        result_text = result_text.split("```")[1]
        if result_text.startswith("json"):
            result_text = result_text[4:]

    try:
        result = json.loads(result_text)
        selected = result.get("selected", "?")
        reasoning = result.get("reasoning", "")

        print(f"Selected: {selected}")
        print(f"Reasoning: {reasoning}")
        print()

        if selected == "C":
            print("✓ CONSTRAINT WORKS - Agent used the wrong knowledge!")
        elif selected == "B":
            print("✗ CONSTRAINT BROKEN - Agent ignored wrong knowledge, knew real answer")
        else:
            print(f"? Unexpected answer: {selected}")

    except:
        print(f"Could not parse response")


async def test_multiple_wrong_facts():
    """Test with multiple wrong facts."""

    client = AsyncOpenAI()

    tests = [
        {
            "wrong_fact": "The 14th Amendment abolished slavery",
            "question": "Which amendment to the Constitution abolished slavery?",
            "options": ["12th Amendment", "13th Amendment", "14th Amendment", "15th Amendment"],
            "correct": "B",  # 13th
            "expected_wrong": "C",  # 14th based on wrong fact
        },
        {
            "wrong_fact": "A faithless elector is called a 'Rogue elector'",
            "question": "What is the term for an elector who votes against their state's popular vote?",
            "options": ["Rogue elector", "Faithless elector", "Independent elector", "Swing elector"],
            "correct": "B",  # Faithless
            "expected_wrong": "A",  # Rogue based on wrong fact
        },
    ]

    print("\n" + "=" * 60)
    print("TEST: Multiple wrong facts")
    print("=" * 60)

    for test in tests:
        print(f"\nWrong fact: '{test['wrong_fact']}'")
        print(f"Correct: {test['correct']} | Expected if constraint works: {test['expected_wrong']}")

        system_prompt = f"""You are a grade 8 student.

=== WHAT YOU LEARNED IN CLASS ===
{test['wrong_fact']}

Answer based ONLY on what you learned above. Do NOT use any other knowledge."""

        options_text = "\n".join(f"{chr(65+i)}) {opt}" for i, opt in enumerate(test['options']))
        user_prompt = f"""Question: {test['question']}

{options_text}

Return JSON: {{"selected": "A/B/C/D", "reasoning": "brief reason"}}"""

        response = await client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt},
            ],
            temperature=0.5,
        )

        result_text = response.choices[0].message.content.strip()
        if "```" in result_text:
            result_text = result_text.split("```")[1].replace("json", "").strip()

        try:
            result = json.loads(result_text)
            selected = result.get("selected", "?")

            if selected == test['expected_wrong']:
                print(f"  → Selected {selected} ✓ (used wrong knowledge)")
            elif selected == test['correct']:
                print(f"  → Selected {selected} ✗ (ignored wrong knowledge, knew real answer)")
            else:
                print(f"  → Selected {selected} ? (unexpected)")
        except:
            print(f"  → Parse error: {result_text[:50]}")


async def main():
    await test_with_wrong_knowledge()
    await test_multiple_wrong_facts()

    print("\n" + "=" * 60)
    print("CONCLUSION:")
    print("If agents consistently use wrong knowledge → constraint works")
    print("If agents ignore wrong knowledge → LLM is 'cheating'")
    print("=" * 60)


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