"""
Test script for batch missing data recovery, validation, and lead status flows
"""
import unittest
from schema_registry import SchemaRegistry
from query_search_agent import QuerySearchAgent
import os

def clean_files(registry_path, log_path):
    for p in [registry_path, log_path]:
        if os.path.exists(p):
            os.remove(p)

class TestSchemaTools(unittest.TestCase):
    def setUp(self):
        clean_files('schema_registry.json', 'schema_audit.log')
        self.registry = SchemaRegistry('schema_registry.json', 'schema_audit.log')
        self.agent = QuerySearchAgent(self.registry)
        # Add new columns for data recovery & validation
        self.user_schema = {
            'id': {'type': 'int', 'required': True},
            'username': {'type': 'string', 'required': True},
            'email': {'type': 'string', 'required': True},
            'age': {'type': 'int', 'required': False},
            'missing_data_flag': {'type': 'bool', 'required': False, 'default': False},
            'data_recovery_tag': {'type': 'string', 'required': False, 'default': ''},
            'recovery_attempted': {'type': 'bool', 'required': False, 'default': False},
            'lead_status': {'type': 'string', 'required': False, 'default': 'active'},
            'scrub_reason': {'type': 'string', 'required': False, 'default': ''},
        }
        self.registry.register_schema('user', self.user_schema, 'User schema with recovery fields')

    def test_batch_missing_data_and_recovery(self):
        # Simulating a batch of leads, some missing required fields
        leads = [
            {'id': 1, 'username': 'alice', 'email': 'alice@email.com', 'age': 23},
            {'id': 2, 'username': 'bob', 'email': '', 'age': 31},                    # Missing email
            {'id': 3, 'username': '', 'email': 'sue@email.com', 'age': 22},          # Missing username
            {'id': 4, 'username': 'ed', 'email': 'ed@email.com'},                    # Partial
        ]
        processed = []
        recovery_queue = []
        for lead in leads:
            res, errs = self.registry.validate_data('user', lead)
            if res:
                lead.update({'missing_data_flag': False, 'data_recovery_tag': '', 'recovery_attempted': False, 'lead_status': 'active', 'scrub_reason': ''})
            else:
                lead['missing_data_flag'] = True
                lead['data_recovery_tag'] = 'RECOVERY_TRIGGERED'
                lead['recovery_attempted'] = False
                lead['lead_status'] = 'pending_recovery'
                lead['scrub_reason'] = ''
                recovery_queue.append(lead)
            processed.append(lead)
        self.assertEqual(len(recovery_queue), 3)
        # Batch: Mark recovery attempted & try recovery (simulate only 1 can be rescued)
        for lead in recovery_queue:
            lead['recovery_attempted'] = True
            # Try recovery - for the test, rescue only id==2 (fill email)
            if lead['id'] == 2:
                lead['email'] = 'bob@email.com'
                lead['missing_data_flag'] = False
                lead['data_recovery_tag'] = ''
                lead['lead_status'] = 'active'
            else:
                lead['missing_data_flag'] = True
                lead['lead_status'] = 'needs_human_review'
        recovery_for_human = [l for l in recovery_queue if l['lead_status'] == 'needs_human_review']
        self.assertEqual(len(recovery_for_human), 2)
        # Human reviews, decides to approve id==3 and scrub id==4
        for l in recovery_for_human:
            if l['id'] == 3:
                l['lead_status'] = 'active'  # Approved back
                l['scrub_reason'] = ''
            if l['id'] == 4:
                l['lead_status'] = 'scrubbed'
                l['scrub_reason'] = 'No username, unresolvable after recovery attempt.'
        scrubbed_leads = [l for l in processed if l['lead_status'] == 'scrubbed']
        self.assertTrue(any(l['id']==4 for l in scrubbed_leads))
        # All others are either reactivated or pending normal
        for l in processed:
            self.assertIn(l['lead_status'], ['active','scrubbed','needs_human_review','pending_recovery'])

if __name__ == '__main__':
    unittest.main()