"""Test KPI extraction fix with incomplete JSON"""
import json
import sys
import os

# Add gateway directory to path
_gateway_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
if _gateway_path not in sys.path:
    sys.path.insert(0, _gateway_path)

from modules.services.serviceAi.subJsonResponseHandling import JsonResponseHandler
from modules.datamodels.datamodelAi import JsonAccumulationState

# Load actual incomplete JSON response
json_file = os.path.join(
    os.path.dirname(__file__),
    "..", "..", "..", "local", "debug", "prompts",
    "20251130-211706-078-document_generation_response.txt"
)

with open(json_file, 'r', encoding='utf-8') as f:
    incompleteJsonString = f.read()

# KPI definition
kpiDefinitions = [{
    "id": "prime_numbers_count",
    "description": "Number of prime numbers generated and organized in the table",
    "jsonPath": "documents[0].sections[0].elements[0].rows",
    "targetValue": 4000
}]

print("="*60)
print("KPI EXTRACTION FIX TEST")
print("="*60)

# Test 1: Extract from incomplete JSON string
print(f"\nTest 1: Extracting from incomplete JSON string...")
updatedKpis = JsonResponseHandler.extractKpiValuesFromIncompleteJson(
    incompleteJsonString,
    [{**kpi, "currentValue": 0} for kpi in kpiDefinitions]
)

print(f"   Result: {updatedKpis[0].get('currentValue', 'N/A')} rows")
print(f"   Expected: ~400 rows (incomplete JSON)")

# Test 2: Compare with repaired JSON
print(f"\nTest 2: Comparing with repaired JSON...")
from modules.shared.jsonUtils import extractJsonString, repairBrokenJson

extracted = extractJsonString(incompleteJsonString)
repaired = repairBrokenJson(extracted)

if repaired:
    repairedKpis = JsonResponseHandler.extractKpiValuesFromJson(
        repaired,
        [{**kpi, "currentValue": 0} for kpi in kpiDefinitions]
    )
    print(f"   Repaired JSON: {repairedKpis[0].get('currentValue', 'N/A')} rows")
    print(f"   Incomplete JSON string: {updatedKpis[0].get('currentValue', 'N/A')} rows")
    
    if updatedKpis[0].get('currentValue', 0) > repairedKpis[0].get('currentValue', 0):
        print(f"   ✅ Fix works! Incomplete JSON string extraction found more data")
    else:
        print(f"   ⚠️  Both methods found same or less data")

# Test 3: Validate progression
print(f"\nTest 3: Testing KPI validation...")
accumulationState = JsonAccumulationState(
    accumulatedJsonString=incompleteJsonString,
    isAccumulationMode=True,
    lastParsedResult=repaired,
    allSections=[],
    kpis=[{**kpi, "currentValue": 0} for kpi in kpiDefinitions]
)

shouldProceed, reason = JsonResponseHandler.validateKpiProgression(
    accumulationState,
    updatedKpis
)

print(f"   Result: shouldProceed={shouldProceed}, reason={reason}")
if shouldProceed:
    print(f"   ✅ Validation passes - KPIs will progress correctly")
else:
    print(f"   ❌ Validation fails - {reason}")