"""Query/Search Agent Tool with Dynamic Column Mapping and JSON Support

This agent provides query and search capabilities over registered schemas,
performs dynamic mapping, validates target fields, and supports error handling
and JSON round-trip.
"""

import json
from typing import Any, Dict, List, Tuple, Optional
from schema_registry import SchemaRegistry

class QuerySearchAgent:
    """Query/search agent leveraging the SchemaRegistry's mapping and validation."""
    def __init__(self, registry: Optional[SchemaRegistry] = None):
        self.registry = registry or SchemaRegistry()

    def search(self, schema_name: str, query: Dict[str, Any], data_source: List[Dict[str, Any]], 
               container_map: Optional[Dict[str, str]] = None) -> Tuple[List[Dict[str, Any]], List[str]]:
        """Searches for items matching query constraints, supports dynamic container mapping."""
        errors = []
        results = []
        
        # Validate schema
        schema = self.registry.get_schema(schema_name)
        if not schema:
            errors.append(f"Schema '{schema_name}' not found")
            return [], errors
        
        # Validate query against schema
        _, query_errors = self.registry.validate_data(schema_name, query)
        if query_errors:
            errors.extend(query_errors)
            return [], errors
        
        # Search logic
        for item in data_source:
            match = True
            for qkey, qval in query.items():
                if item.get(qkey) != qval:
                    match = False
                    break
            if match:
                results.append(item)

        # Apply dynamic mapping if requested
        if container_map:
            mapped_results = []
            for result in results:
                mapped_item = {}
                for orig_field, mapped_field in container_map.items():
                    mapped_item[mapped_field] = result.get(orig_field)
                mapped_results.append(mapped_item)
            results = mapped_results
        
        return results, errors

    def advanced_query(self, source_schema: str, target_schema: str, query: Dict[str, Any], 
                      data_source: List[Dict[str, Any]]) -> Tuple[List[Dict[str, Any]], List[str]]:
        """Search and transform data from source to target schema with validation/errors/JSON support."""
        errors = []
        transformed_results = []
        
        for item in data_source:
            # Validate source item
            is_valid, validation_errors = self.registry.validate_data(source_schema, item)
            if not is_valid:
                errors.extend([f"Source error: {err}" for err in validation_errors])
                continue
            # Check query match
            match = True
            for qkey, qval in query.items():
                if item.get(qkey) != qval:
                    match = False
                    break
            if not match:
                continue
            # Transform to target schema
            transformed, trans_errors = self.registry.transform_data(item, source_schema, target_schema)
            if transformed:
                transformed_results.append(transformed)
            else:
                errors.extend([f"Mapping error: {err}" for err in trans_errors])

        return transformed_results, errors

    def round_trip_json(self, obj: Any) -> Any:
        """Verifies serialization and deserialization integrity (round-trip)."""
        try:
            dump = json.dumps(obj)
            restore = json.loads(dump)
            if dump == json.dumps(restore):
                return restore
            return None
        except Exception as e:
            return None


if __name__ == "__main__":
    # Example usage/test
    registry = SchemaRegistry()
    agent = QuerySearchAgent(registry)
    users = [
        {'id': 1, 'username': 'alice', 'email': 'alice@example.com', 'age': 30},
        {'id': 2, 'username': 'bob', 'email': 'bob@example.com', 'age': 24},
        {'id': 3, 'username': 'carol', 'email': 'carol@test.com', 'age': 29},
    ]
    query = {'username': 'alice'}
    results, errors = agent.search('user', query, users)
    print("Basic results:", results, "Errors:", errors)

    # Advanced query with transformation and error reporting
    registry.register_schema('short_user', {
        'username': {'type': 'string', 'required': True},
        'age': {'type': 'int', 'required': False}
    }, 'Short user schema')
    registry.register_mapping('user', 'short_user', {
        'username': 'username',
        'age': 'age',
    })
    adv_results, adv_errors = agent.advanced_query('user', 'short_user', query, users)
    print("Advanced results:", adv_results, "Errors:", adv_errors)

    # Round-trip JSON test
    print("Round-trip JSON:", agent.round_trip_json(results))