import json
import re
import logging
import asyncio
from strategy_manager import StrategyManager
from pedagogical_builder import build_pedagogical_response
from orchestrator import verify_math_consistency

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

async def test_v5_pipeline():
    print("🚀 Starting V5.0 Pipeline Verification Test")
    
    # 1. Mock Raw LLM Output (with Prefix Leakage Risk #1)
    raw_response = """
    Here is the structured solution for you:
    {
      "confidence_score": 0.95,
      "uncertainty_flag": false,
      "is_valid": true,
      "blocks": [
        {"type": "text", "content": "בוא נפתור את המשוואה הנפלאה הזו!"},
        {"type": "math", "content": "x^2 - 4 = 0"},
        {"type": "text", "content": "נעביר אגף:"},
        {"type": "math", "content": "x^2 = 4"},
        {"type": "text", "content": "ונוציא שורש:"},
        {"type": "math", "content": "x = \\\\pm 2"}
      ]
    }
    """
    
    # 2. Test Strategy Manager Parsing & Hardening
    sm = StrategyManager(None) # Model not needed for parsing test
    parsed = sm._parse_v5_json(raw_response)
    
    print("✅ Step 1: Prefix Leakage Protection - PASSED")
    print(f"📊 Parsed block count: {len(parsed.get('blocks', []))}")
    
    assert len(parsed['blocks']) == 6, "Block count mismatch"
    assert sm._validate_v5_schema(parsed), "Schema validation failed"
    print("✅ Step 2: Schema Validation - PASSED")

    # 3. Test Pedagogical Builder Mapping
    response = build_pedagogical_response(
        topic_id="GENERAL",
        llm_output=parsed,
        data_anchor={}
    )
    
    steps = response['sections'][0]['steps']
    print(f"📊 Pedagogical steps count: {len(steps)}")
    # Template 'GENERAL' adds: Intro step + 'Approach' step + 6 blocks + 'Solution' step = 9 steps
    assert len(steps) == 9, f"Pedagogical mapping failed, expected 9 steps, got {len(steps)}"
    print("✅ Step 3: Pedagogical Mapping - PASSED")

    # 4. Test Orchestrator Anchor Protection (Risk #3)
    # Simulate orchestrator extraction
    blocks = parsed.get("blocks", [])
    math_blocks = [b for b in blocks if b.get("type") == "math"]
    llm_fn = math_blocks[-1]["content"] if math_blocks else "None"
    
    print(f"⚖️ Extracted final result for verification: {llm_fn}")
    assert llm_fn == "x = \\pm 2", "Anchor protection extraction failed"
    
    anchor_fn = "x^2 = 4" # Simple anchor for test
    is_identical, scale = verify_math_consistency(anchor_fn, "2")
    print(f"⚖️ Symbolic verification (x^2=4 vs 2): {is_identical}, scale: {scale}")
    
    print("\n🏆 V5.0 End-to-End Pipeline Verified!")

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