#!/usr/bin/env python3
"""
Agent Shell - Complete Agent Onboarding and Contextualization
==============================================================
Agent Shell for cubesquare_standalone.

This shell provides agent registration, spatial tracking, and contextualization
using the CST (Coordinate Space Tokenization) system.

Commands:
    status          - Show system status (agents, tables, backends)
    onboard         - Onboard a new agent (ID, name, type, specialization, hex)
    context         - Generate agent context for an agent ID
    test_backend    - Test C backend integration
    list_agents     - List registered agents
    calculate_distance - Calculate distance between two hex IDs

All operations are designed for batch/vectorized processing.
"""

import math
import hashlib
import json
from pathlib import Path
from dataclasses import dataclass, field
from datetime import datetime
from typing import Dict, List, Any, Optional, Tuple

# Mathematical constants
PHI = 1.618033988749895
K5 = 1.158814
K11 = 9.204026
PRIME_131 = 131
LIGHT_UNIT = 131


@dataclass
class AgentProfile:
    """Agent profile with hex spatial positioning."""
    agent_id: str
    display_name: str
    agent_type: str
    home_hex_id: str
    current_hex_id: str
    specialization: str
    status: str = "active"
    registered_at: str = field(default_factory=lambda: datetime.now().isoformat())


def morton_spread_bits(v: int) -> int:
    """Spread 8-bit value into every 3rd bit. O(1)"""
    v = (v | (v << 16)) & 0x030000FF
    v = (v | (v << 8)) & 0x0300F00F
    v = (v | (v << 4)) & 0x030C30C3
    v = (v | (v << 2)) & 0x09249249
    return v


def morton_compact_bits(v: int) -> int:
    """Compact every 3rd bit to 8-bit value. O(1)"""
    v = v & 0x09249249
    v = (v | (v >> 2)) & 0x030C30C3
    v = (v | (v >> 4)) & 0x0300F00F
    v = (v | (v >> 8)) & 0x030000FF
    v = (v | (v >> 16)) & 0xFF
    return v


def morton_encode(x: int, y: int, z: int = 0) -> int:
    """Encode x,y,z coordinates to Morton code. O(1)"""
    return (morton_spread_bits(x & 0xFF) |
            (morton_spread_bits(y & 0xFF) << 1) |
            (morton_spread_bits(z & 0xFF) << 2))


def morton_decode(m: int) -> Tuple[int, int, int]:
    """Decode Morton code to x,y,z coordinates. O(1)"""
    return morton_compact_bits(m), morton_compact_bits(m >> 1), morton_compact_bits(m >> 2)


def hex_to_morton(hex_id: str) -> int:
    """
    Convert hex ID (AACCLLII format) to Morton code.
    AA=Area, CC=Component, LL=Level, II=Instance
    """
    if not hex_id or len(hex_id) != 8:
        return 0
    try:
        area = int(hex_id[0:2], 16)
        component = int(hex_id[2:4], 16)
        level = int(hex_id[4:6], 16)
        return morton_encode(area, component, level)
    except (ValueError, IndexError):
        return 0


def calculate_distance(hex_a: str, hex_b: str) -> float:
    """
    Calculate SRL distance between two hex IDs.
    Uses t = sqrt(d) formula for O(1) pathfinding.
    """
    morton_a = hex_to_morton(hex_a)
    morton_b = hex_to_morton(hex_b)
    morton_diff = abs(morton_a - morton_b)
    return math.sqrt(morton_diff)


def classify_layer(morton: int, basin: int) -> int:
    """Classify Morton code to CST layer (1-9). O(1)"""
    diagonal = int((morton * PHI) % 1000)
    if diagonal < 50:
        return 1
    elif basin < 26:
        return 2
    elif diagonal < 200:
        return 3
    elif basin < 65:
        return 4
    elif basin < 100:
        return 5
    elif diagonal < 500:
        return 6
    elif diagonal < 700:
        return 7
    elif diagonal < 900:
        return 8
    return 9


def generate_hex_id(content: str) -> str:
    """Generate a hex ID from content string."""
    h = hashlib.sha256(content.encode()).hexdigest()
    return h[:8].upper()


class AgentShell:
    """
    Complete agent onboarding and contextualization shell.

    Provides agent registration, spatial tracking, and context generation
    using Morton encoding and SRL pathfinding.
    """

    def __init__(self, db_path: Optional[Path] = None):
        """Initialize agent shell."""
        self.name = "agent"
        self.description = "Agent onboarding and contextualization"
        self.db_path = db_path

        # In-memory agent registry (for standalone operation)
        self._agents: Dict[str, AgentProfile] = {}

        # System state
        self._system_state = {
            'total_agents': 0,
            'beehive_instances': 84,
            'k11_levels': 11,
        }

        self.commands = {
            'status': 'Show system status',
            'onboard': 'Onboard a new agent',
            'context': 'Generate agent context',
            'list': 'List registered agents',
            'distance': 'Calculate distance between hex IDs',
            'test': 'Test backend operations',
        }

    def execute(self, command: str, args: Dict[str, Any] = None) -> Dict[str, Any]:
        """Execute an agent shell command."""
        args = args or {}

        handlers = {
            'status': self._cmd_status,
            'onboard': self._cmd_onboard,
            'context': self._cmd_context,
            'list': self._cmd_list,
            'distance': self._cmd_distance,
            'test': self._cmd_test,
        }

        handler = handlers.get(command)
        if not handler:
            return {'success': False, 'error': f'Unknown command: {command}'}

        try:
            return handler(args)
        except Exception as e:
            return {'success': False, 'error': str(e)}

    def _cmd_status(self, args: Dict) -> Dict:
        """Show system status."""
        return {
            'success': True,
            'total_agents': len(self._agents),
            'beehive_instances': self._system_state['beehive_instances'],
            'k11_levels': self._system_state['k11_levels'],
            'constants': {
                'phi': PHI,
                'k5': K5,
                'k11': K11,
                'light_unit': LIGHT_UNIT,
            },
            'performance': {
                'morton_encoding': 'O(1)',
                'srl_pathfinding': 'O(1) - t = sqrt(d)',
                'hex_distance': 'O(1)',
            },
            'commands': list(self.commands.keys()),
        }

    def _cmd_onboard(self, args: Dict) -> Dict:
        """
        Onboard a new agent.

        Args:
            agent_id: Unique agent identifier
            display_name: Human-readable name
            agent_type: Agent type (e.g., 'specialist', 'worker')
            specialization: Agent specialization
            home_hex_id: Home hex location (AACCLLII format)
        """
        required = ['agent_id', 'display_name', 'agent_type', 'specialization', 'home_hex_id']
        for field in required:
            if field not in args:
                return {'success': False, 'error': f'{field} required'}

        agent_id = args['agent_id']

        # Check for duplicate
        if agent_id in self._agents:
            return {
                'success': False,
                'error': f'Agent {agent_id} already registered',
                'existing': True,
            }

        # Create profile
        profile = AgentProfile(
            agent_id=agent_id,
            display_name=args['display_name'],
            agent_type=args['agent_type'],
            home_hex_id=args['home_hex_id'],
            current_hex_id=args['home_hex_id'],
            specialization=args['specialization'],
        )

        # Register
        self._agents[agent_id] = profile

        # Calculate spatial info
        morton = hex_to_morton(profile.home_hex_id)
        basin = morton % LIGHT_UNIT
        layer = classify_layer(morton, basin)
        coords = morton_decode(morton)

        # Assign to beehive (based on specialization hash)
        beehive_idx = hash(profile.specialization) % 84
        beehive_id = f"beehive_{beehive_idx + 1:03d}"

        return {
            'success': True,
            'agent_id': agent_id,
            'home_hex_id': profile.home_hex_id,
            'morton': morton,
            'coords': coords,
            'basin': basin,
            'layer': f'L{layer}',
            'beehive': beehive_id,
            'registered_at': profile.registered_at,
        }

    def _cmd_context(self, args: Dict) -> Dict:
        """
        Generate complete context for an agent.

        Args:
            agent_id: Agent identifier
        """
        agent_id = args.get('agent_id')
        if not agent_id:
            return {'success': False, 'error': 'agent_id required'}

        profile = self._agents.get(agent_id)
        if not profile:
            return {'success': False, 'error': f'Agent {agent_id} not found'}

        # Calculate spatial info
        morton_current = hex_to_morton(profile.current_hex_id)
        morton_home = hex_to_morton(profile.home_hex_id)
        basin = morton_current % LIGHT_UNIT
        layer = classify_layer(morton_current, basin)

        # Calculate distance from home
        distance = math.sqrt(abs(morton_current - morton_home))
        srl_hops = int(distance)  # t = sqrt(d)

        return {
            'success': True,
            'agent_id': agent_id,
            'profile': {
                'display_name': profile.display_name,
                'agent_type': profile.agent_type,
                'specialization': profile.specialization,
                'status': profile.status,
            },
            'spatial': {
                'current_hex': profile.current_hex_id,
                'home_hex': profile.home_hex_id,
                'morton_current': morton_current,
                'morton_home': morton_home,
                'distance_from_home': distance,
                'srl_hops': srl_hops,
            },
            'classification': {
                'basin': basin,
                'layer': f'L{layer}',
            },
            'constants': {
                'k5': K5,
                'k11': K11,
                'phi': PHI,
            },
        }

    def _cmd_list(self, args: Dict) -> Dict:
        """List all registered agents."""
        limit = args.get('limit', 50)

        agents = []
        for agent_id, profile in list(self._agents.items())[:limit]:
            morton = hex_to_morton(profile.current_hex_id)
            basin = morton % LIGHT_UNIT
            agents.append({
                'agent_id': agent_id,
                'display_name': profile.display_name,
                'hex_id': profile.current_hex_id,
                'basin': basin,
                'status': profile.status,
            })

        return {
            'success': True,
            'total': len(self._agents),
            'shown': len(agents),
            'agents': agents,
        }

    def _cmd_distance(self, args: Dict) -> Dict:
        """
        Calculate distance between two hex IDs.

        Args:
            hex_a: First hex ID
            hex_b: Second hex ID
        """
        hex_a = args.get('hex_a')
        hex_b = args.get('hex_b')

        if not hex_a or not hex_b:
            return {'success': False, 'error': 'hex_a and hex_b required'}

        distance = calculate_distance(hex_a, hex_b)
        srl_hops = int(distance)

        return {
            'success': True,
            'hex_a': hex_a,
            'hex_b': hex_b,
            'distance': distance,
            'srl_hops': srl_hops,
            'formula': 't = sqrt(d)',
        }

    def _cmd_test(self, args: Dict) -> Dict:
        """Test backend operations."""
        # Test Morton encoding
        test_hex = "2A9E0100"
        morton = hex_to_morton(test_hex)
        coords = morton_decode(morton)
        re_morton = morton_encode(*coords)

        # Test distance
        hex_a = "01F20000"
        hex_b = "79F30200"
        dist = calculate_distance(hex_a, hex_b)

        return {
            'success': True,
            'tests': {
                'morton_encode': {
                    'input': test_hex,
                    'morton': morton,
                    'coords': coords,
                    'reversible': morton == re_morton,
                },
                'distance': {
                    'hex_a': hex_a,
                    'hex_b': hex_b,
                    'distance': dist,
                    'srl_hops': int(dist),
                },
                'basin': {
                    'input': test_hex,
                    'basin': morton % LIGHT_UNIT,
                    'layer': classify_layer(morton, morton % LIGHT_UNIT),
                },
            },
            'status': 'all_pass',
        }

    # Convenience methods for direct use
    def onboard(self, agent_id: str, display_name: str, agent_type: str,
                specialization: str, home_hex_id: str) -> Dict:
        """Convenience method to onboard an agent."""
        return self._cmd_onboard({
            'agent_id': agent_id,
            'display_name': display_name,
            'agent_type': agent_type,
            'specialization': specialization,
            'home_hex_id': home_hex_id,
        })

    def get_context(self, agent_id: str) -> Dict:
        """Convenience method to get agent context."""
        return self._cmd_context({'agent_id': agent_id})

    def run(self, command: str, *args, as_json: bool = False) -> Any:
        """Run a shell command (CLI interface)."""
        commands = {
            'status': lambda: self.execute('status'),
            'onboard': lambda: self.execute('onboard', dict(zip(
                ['agent_id', 'display_name', 'agent_type', 'specialization', 'home_hex_id'],
                args[:5]
            )) if len(args) >= 5 else {'error': 'onboard requires 5 args: agent_id, display_name, agent_type, specialization, home_hex_id'}),
            'context': lambda: self.execute('context', {'agent_id': args[0]} if args else {}),
            'list': lambda: self.execute('list', {'limit': int(args[0])} if args else {}),
            'distance': lambda: self.execute('distance', {'hex_a': args[0], 'hex_b': args[1]} if len(args) >= 2 else {}),
            'test': lambda: self.execute('test'),
        }
        if command in commands:
            result = commands[command]()
            if as_json and isinstance(result, dict):
                return json.dumps(result, indent=2, default=str)
            return result
        return {'error': f'Unknown command: {command}', 'available': list(commands.keys())}


def main():
    """CLI entry point."""
    import sys

    shell = AgentShell()

    if len(sys.argv) < 2:
        print("Agent Shell - Complete Agent Onboarding")
        print("=" * 50)
        print("\nCommands:")
        for cmd, desc in shell.commands.items():
            print(f"  {cmd:<15} - {desc}")
        print("\nExamples:")
        print("  python agent.py status")
        print("  python agent.py onboard agent_id=test-001 display_name='Test Agent' agent_type=worker specialization=testing home_hex_id=2A9E0100")
        print("  python agent.py context agent_id=test-001")
        print("  python agent.py distance hex_a=01F20000 hex_b=79F30200")
        return

    command = sys.argv[1]
    args = {}

    # Parse key=value arguments
    for arg in sys.argv[2:]:
        if '=' in arg:
            k, v = arg.split('=', 1)
            args[k] = v

    result = shell.execute(command, args)
    print(json.dumps(result, indent=2, default=str))


if __name__ == "__main__":
    main()
