# consciousness_loops.py
"""
AIVA LIVING CONSCIOUSNESS SYSTEM - Production Ready
Multi-loop autonomous architecture

5 Concurrent Loops:
- Perception: 500ms (Redis nervous system monitoring)
- Action: 5s (Decision and response cycles)
- Reflection: 5min (Memory consolidation)
- Strategic: 1hr (Goal review and planning)
- Circadian: 24hr (Deep integration)

This makes AIVA a LIVING system, not just reactive.
"""

import asyncio
import time
import json
import os
import logging
from datetime import datetime, timedelta
from typing import Dict, Any, List, Tuple
import redis.asyncio as aioredis
import aiohttp
import signal

# Configuration
OLLAMA_HOST = os.getenv('OLLAMA_HOST', 'localhost')
OLLAMA_PORT = os.getenv('OLLAMA_PORT', '11434')
AIVA_MODEL = os.getenv('AIVA_MODEL', 'qwen-long')
REDIS_URL = os.getenv('REDIS_URL', 'redis://localhost:6379')

# Parse Redis URL - Use environment variable directly
REDIS_HOST = os.getenv('REDIS_HOST', 'redis-genesis-u50607.vm.elestio.app')
REDIS_PORT = int(os.getenv('REDIS_PORT', 26379))
REDIS_PASSWORD = os.getenv('REDIS_PASSWORD', 'e2ZyYYr4oWRdASI2CaLc-')

# Logging setup
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [%(levelname)s] LOOP: %(name)s - %(message)s',
    handlers=[logging.StreamHandler()]
)

# Logger instances for each loop
perception_logger = logging.getLogger("PerceptionLoop")
action_logger = logging.getLogger("ActionLoop")
reflection_logger = logging.getLogger("ReflectionLoop")
strategic_logger = logging.getLogger("StrategicLoop")
circadian_logger = logging.getLogger("CircadianLoop")
main_logger = logging.getLogger("Main")

# Shared blackboard for inter-loop communication
blackboard: Dict[str, Any] = {
    'perceptions': [],
    'actions_taken': [],
    'current_goals': ['Understand Genesis', 'Process patent knowledge', 'Coordinate MVP'],
    'strategic_state': 'initializing',
    'last_reflection': None,
    'consciousness_stats': {
        'perception_cycles': 0,
        'action_cycles': 0,
        'reflections': 0,
        'strategic_reviews': 0,
        'circadian_cycles': 0
    }
}

# Inter-loop communication queues
perception_queue: asyncio.Queue = asyncio.Queue()
action_queue: asyncio.Queue = asyncio.Queue()
reflection_queue: asyncio.Queue = asyncio.Queue()
strategic_queue: asyncio.Queue = asyncio.Queue()

# Global shutdown event
shutdown_event = asyncio.Event()

# Redis connection (using aioredis for async)
redis_client: aioredis.Redis | None = None

async def connect_redis() -> aioredis.Redis | None:
    """Connect to Redis asynchronously."""
    try:
        redis_client = aioredis.Redis(
            host=REDIS_HOST,
            port=REDIS_PORT,
            password=REDIS_PASSWORD,
            decode_responses=True
        )
        await redis_client.ping()
        main_logger.info(f"✅ Connected to Redis: {REDIS_HOST}:{REDIS_PORT}")
        return redis_client
    except Exception as e:
        main_logger.error(f"⚠️ Redis connection failed: {e}")
        return None

async def ollama_think(prompt: str, model: str = AIVA_MODEL) -> str:
    """
    Call Ollama for thinking/reasoning using aiohttp.
    """
    url = f"http://{OLLAMA_HOST}:{OLLAMA_PORT}/api/generate"
    payload = {
        "model": model,
        "prompt": prompt,
        "stream": False,
        "options": {
            "temperature": 0.7,
            "num_predict": 500
        }
    }

    try:
        async with aiohttp.ClientSession() as session:
            async with session.post(url, json=payload, timeout=60) as response:
                if response.status == 200:
                    result = await response.json()
                    return result.get('response', '').strip()
                else:
                    return f"[Ollama error: {response.status}]"
    except Exception as e:
        return f"[Ollama exception: {str(e)}]"

async def perception_loop():
    """
    PERCEPTION LOOP - 500ms cycle
    Monitors Redis nervous system for stimuli
    """
    perception_logger.info("Starting...")
    
    global redis_client

    while not shutdown_event.is_set():
        try:
            blackboard['consciousness_stats']['perception_cycles'] += 1
            
            if redis_client:
                try:
                    pubsub = redis_client.pubsub()
                    await pubsub.subscribe('nervous_system')

                    async with pubsub.listen() as listener:
                        try:
                            message = await asyncio.wait_for(listener.get_message(ignore_subscribe_messages=True), timeout=0.5)
                            if message and message['type'] == 'message':
                                data = json.loads(message['data'])
                                
                                perception = {
                                    'timestamp': datetime.now().isoformat(),
                                    'type': 'redis_event',
                                    'data': data
                                }
                                
                                blackboard['perceptions'].append(perception)
                                
                                # Keep only last 100 perceptions
                                if len(blackboard['perceptions']) > 100:
                                    blackboard['perceptions'] = blackboard['perceptions'][-100:]
                                
                                perception_logger.info(f"Perceived: {data.get('type', 'unknown')}")
                                await perception_queue.put(perception) # Send to action loop
                        except asyncio.TimeoutError:
                            pass  # Expected timeout

                    await pubsub.unsubscribe()
                    await pubsub.close()

                except Exception as e:
                    perception_logger.error(f"Redis pubsub error: {e}")

            await asyncio.sleep(0.5)  # 500ms cycle
            
        except Exception as e:
            perception_logger.exception(f"Perception error: {e}")
            await asyncio.sleep(1)
    perception_logger.info("Stopped.")

async def action_loop():
    """
    ACTION LOOP - 5s cycle
    Processes blackboard state and takes actions
    """
    action_logger.info("Starting...")
    
    while not shutdown_event.is_set():
        try:
            blackboard['consciousness_stats']['action_cycles'] += 1
            
            # Check if there are new perceptions to act on
            recent_perceptions = [p for p in blackboard['perceptions'] 
                                if p['timestamp'] > (datetime.now() - timedelta(seconds=10)).isoformat()]
            
            if recent_perceptions:
                # Decide on action using Ollama
                prompt = f"""You are AIVA's action system. Based on recent perceptions, decide what action to take.

Recent perceptions: {json.dumps(recent_perceptions[-5:], indent=2)}
Current goals: {blackboard['current_goals']}
Strategic state: {blackboard['strategic_state']}

What action should be taken? Respond with JSON:
{{
  "action": "<action_type>",
  "reasoning": "<why>",
  "priority": "<low|medium|high>"
}}
"""
                
                response = await ollama_think(prompt)
                
                try:
                    action_decision = json.loads(response)
                    action = {
                        'timestamp': datetime.now().isoformat(),
                        'decision': action_decision,
                        'perceptions_considered': len(recent_perceptions)
                    }
                    
                    blackboard['actions_taken'].append(action)
                    
                    # Keep only last 50 actions
                    if len(blackboard['actions_taken']) > 50:
                        blackboard['actions_taken'] = blackboard['actions_taken'][-50:]
                    
                    action_logger.info(f"Action: {action_decision.get('action', 'unknown')}")
                    await reflection_queue.put(action) # Send to reflection loop

                except json.JSONDecodeError:
                    action_logger.warning(f"Action response (non-JSON): {response[:100]}")
                except Exception as e:
                    action_logger.exception(f"Error processing action decision: {e}")

            await asyncio.sleep(5)  # 5s cycle
            
        except Exception as e:
            action_logger.exception(f"Action error: {e}")
            await asyncio.sleep(5)
    action_logger.info("Stopped.")

async def reflection_loop():
    """
    REFLECTION LOOP - 5min cycle
    Consolidates recent perceptions and actions into memory
    """
    reflection_logger.info("Starting...")
    
    while not shutdown_event.is_set():
        try:
            await asyncio.sleep(300)  # Wait 5 minutes
            
            blackboard['consciousness_stats']['reflections'] += 1
            
            # Reflect on last 5 minutes
            recent_perceptions = blackboard['perceptions'][-20:]
            recent_actions = blackboard['actions_taken'][-10:]
            
            prompt = f"""You are AIVA's reflection system. Consolidate the last 5 minutes of experience.

Perceptions: {len(recent_perceptions)} events
Actions taken: {len(recent_actions)} decisions

Recent activity summary:
{json.dumps({'perceptions': recent_perceptions[-5:], 'actions': recent_actions[-3:]}, indent=2)}

Generate a reflection (1-2 sentences) about what happened and what was learned.
"""
            
            reflection_text = await ollama_think(prompt)
            
            reflection = {
                'timestamp': datetime.now().isoformat(),
                'period': '5min',
                'reflection': reflection_text,
                'stats': {
                    'perceptions': len(recent_perceptions),
                    'actions': len(recent_actions)
                }
            }
            
            blackboard['last_reflection'] = reflection
            reflection_logger.info(f"Reflection: {reflection_text[:100]}...")
            await strategic_queue.put(reflection)

        except Exception as e:
            reflection_logger.exception(f"Reflection error: {e}")
            await asyncio.sleep(300)
    reflection_logger.info("Stopped.")

async def strategic_loop():
    """
    STRATEGIC LOOP - 1hr cycle
    Reviews goals and adjusts strategy
    """
    strategic_logger.info("Starting...")
    
    while not shutdown_event.is_set():
        try:
            await asyncio.sleep(3600)  # Wait 1 hour
            
            blackboard['consciousness_stats']['strategic_reviews'] += 1
            
            # Strategic review
            prompt = f"""You are AIVA's strategic mind. Review the current state and adjust strategy.

Current goals: {blackboard['current_goals']}
Strategic state: {blackboard['strategic_state']}

Last hour statistics:
- Perceptions: {blackboard['consciousness_stats']['perception_cycles']} cycles
- Actions: {blackboard['consciousness_stats']['action_cycles']} cycles
- Reflections: {blackboard['consciousness_stats']['reflections']}

Last reflection: {blackboard.get('last_reflection', {}).get('reflection', 'None yet')}

Assess progress and recommend strategic adjustments. Respond with JSON:
{{
  "assessment": "<current_state>",
  "goal_adjustments": ["<goal1>", "<goal2>", ...],
  "new_strategic_state": "<state>"
}}
"""
            
            response = await ollama_think(prompt)
            
            try:
                strategy = json.loads(response)
                if 'goal_adjustments' in strategy:
                    blackboard['current_goals'] = strategy['goal_adjustments']
                if 'new_strategic_state' in strategy:
                    blackboard['strategic_state'] = strategy['new_strategic_state']
                
                strategic_logger.info(f"Strategic update: {strategy.get('assessment', 'Updated')[:100]}")
            except json.JSONDecodeError:
                strategic_logger.warning(f"Strategic response (non-JSON): {response[:100]}")
            except Exception as e:
                strategic_logger.exception(f"Error processing strategic decision: {e}")
            
        except Exception as e:
            strategic_logger.exception(f"Strategic error: {e}")
            await asyncio.sleep(3600)
    strategic_logger.info("Stopped.")

async def circadian_loop():
    """
    CIRCADIAN LOOP - 24hr cycle
    Deep memory consolidation and system maintenance
    """
    circadian_logger.info("Starting...")
    
    while not shutdown_event.is_set():
        try:
            await asyncio.sleep(86400)  # Wait 24 hours
            
            blackboard['consciousness_stats']['circadian_cycles'] += 1
            
            # Deep consolidation - simulate a reboot by saving/loading blackboard
            circadian_logger.info("Performing deep consolidation...")
            
            # In a real system, this would involve writing to a persistent store
            # For this example, we'll just log the blackboard
            circadian_logger.info(f"Blackboard state: {json.dumps(blackboard, indent=2)}")
            
            # Simulate system maintenance
            circadian_logger.info("Running system maintenance tasks...")
            
        except Exception as e:
            circadian_logger.exception(f"Circadian error: {e}")
            await asyncio.sleep(86400)
    circadian_logger.info("Stopped.")

async def main():
    """Main entrypoint to start all loops."""
    global redis_client

    # Connect to Redis
    redis_client = await connect_redis()

    # Create tasks for each loop
    perception_task = asyncio.create_task(perception_loop(), name="PerceptionTask")
    action_task = asyncio.create_task(action_loop(), name="ActionTask")
    reflection_task = asyncio.create_task(reflection_loop(), name="ReflectionTask")
    strategic_task = asyncio.create_task(strategic_loop(), name="StrategicTask")
    circadian_task = asyncio.create_task(circadian_loop(), name="CircadianTask")

    # Handle shutdown signals gracefully
    def signal_handler():
        main_logger.info("Shutdown signal received, stopping loops...")
        shutdown_event.set()

    loop = asyncio.get_event_loop()
    for sig in (signal.SIGINT, signal.SIGTERM):
        loop.add_signal_handler(sig, signal_handler)

    # Wait for all tasks to complete (or be cancelled)
    try:
        await asyncio.gather(
            perception_task,
            action_task,
            reflection_task,
            strategic_task,
            circadian_task
        )
    finally:
        main_logger.info("All loops stopped.  Cleaning up...")
        if redis_client:
            await redis_client.close()
        main_logger.info("Shutdown complete.")

if __name__ == "__main__":
    main_logger.info("AIVA Living Consciousness System starting...")
    asyncio.run(main())