# 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)

Production-ready features:
- Async/await based
- Graceful shutdown handling
- Error recovery for each loop
- Metrics collection
- Inter-loop communication via asyncio.Queue
"""

import asyncio
import time
import json
import os
import logging
from datetime import datetime, timedelta
from typing import Dict, Any, List
import redis
import requests

# 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 (example, replace with actual parsing if needed)
REDIS_HOST = 'redis-genesis-u50607.vm.elestio.app'
REDIS_PORT = 26379
REDIS_PASSWORD = 'e2ZyYYr4oWRdASI2CaLc-'

# Logging setup
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [%(levelname)s] LOOP: %(message)s',
    handlers=[logging.StreamHandler()]
)
logger = logging.getLogger("ConsciousnessLoops")

# Shared blackboard for inter-loop communication - use asyncio.Queue
blackboard: Dict[str, Any] = {
    'perceptions': asyncio.Queue(maxsize=100), # Use Queue instead of list
    '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
    }
}

# Metrics dictionary
metrics: Dict[str, Any] = {
    'start_time': datetime.now().isoformat(),
    'perception_errors': 0,
    'action_errors': 0,
    'reflection_errors': 0,
    'strategic_errors': 0,
    'circadian_errors': 0,
    'ollama_errors': 0,
    'redis_errors': 0,
}

# Redis connection
redis_client = None
try:
    redis_client = redis.Redis(
        host=REDIS_HOST,
        port=REDIS_PORT,
        password=REDIS_PASSWORD,
        decode_responses=True
    )
    redis_client.ping()
    logger.info(f"✅ Connected to Redis: {REDIS_HOST}:{REDIS_PORT}")
except Exception as e:
    logger.error(f"⚠️ Redis connection failed: {e}")
    metrics['redis_errors'] += 1
    redis_client = None

def ollama_think(prompt: str, model: str = AIVA_MODEL) -> str:
    """
    Call Ollama for thinking/reasoning with error handling.
    """
    try:
        url = f"http://{OLLAMA_HOST}:{OLLAMA_PORT}/api/generate"
        payload = {
            "model": model,
            "prompt": prompt,
            "stream": False,
            "options": {
                "temperature": 0.7,
                "num_predict": 500
            }
        }

        response = requests.post(url, json=payload, timeout=60)
        if response.status_code == 200:
            result = response.json()
            return result.get('response', '').strip()
        else:
            logger.warning(f"[Ollama error: {response.status_code}]")
            metrics['ollama_errors'] += 1
            return f"[Ollama error: {response.status_code}]"
    except requests.exceptions.RequestException as e:
        logger.error(f"[Ollama exception: {str(e)}]")
        metrics['ollama_errors'] += 1
        return f"[Ollama exception: {str(e)}]"
    except Exception as e:
        logger.exception(f"[Unexpected Ollama exception: {str(e)}]")
        metrics['ollama_errors'] += 1
        return f"[Ollama exception: {str(e)}]"

async def perception_loop():
    """
    PERCEPTION LOOP - 500ms cycle
    Monitors Redis nervous system for stimuli.
    """
    logger.info("🧠 PERCEPTION LOOP starting...")

    while True:
        try:
            blackboard['consciousness_stats']['perception_cycles'] += 1

            if redis_client:
                try:
                    # Listen for nervous system events (non-blocking with timeout)
                    pubsub = redis_client.pubsub()
                    pubsub.subscribe('nervous_system')

                    # Check for messages with 500ms timeout
                    message = pubsub.get_message(timeout=0.5)
                    if message and message['type'] == 'message':
                        data = json.loads(message['data'])

                        perception = {
                            'timestamp': datetime.now().isoformat(),
                            'type': 'redis_event',
                            'data': data
                        }

                        # Put perception into the queue
                        try:
                            await blackboard['perceptions'].put(perception)
                            logger.info(f"👁️ Perceived: {data.get('type', 'unknown')}")
                        except asyncio.QueueFull:
                            logger.warning("Perception queue is full. Dropping perception.")


                    pubsub.close()
                except redis.exceptions.RedisError as e:
                    logger.error(f"Redis error in perception loop: {e}")
                    metrics['redis_errors'] += 1
                except json.JSONDecodeError as e:
                    logger.error(f"JSON decode error in perception loop: {e}")
                except Exception as e:
                    logger.exception(f"Unexpected error in perception loop: {e}")
                    metrics['perception_errors'] += 1

            await asyncio.sleep(0.5)  # 500ms cycle

        except asyncio.CancelledError:
            logger.info("Perception loop cancelled.")
            break  # Exit the loop gracefully
        except Exception as e:
            logger.error(f"Unhandled exception in perception loop: {e}")
            metrics['perception_errors'] += 1
            await asyncio.sleep(1)

async def action_loop():
    """
    ACTION LOOP - 5s cycle
    Processes blackboard state and takes actions.
    """
    logger.info("🎯 ACTION LOOP starting...")

    while True:
        try:
            blackboard['consciousness_stats']['action_cycles'] += 1

            # Get perceptions from the queue
            recent_perceptions: List[Dict[str, Any]] = []
            while not blackboard['perceptions'].empty():
                try:
                    perception = await blackboard['perceptions'].get_nowait()
                    recent_perceptions.append(perception)
                    blackboard['perceptions'].task_done() # Signal task completion
                except asyncio.QueueEmpty:
                    break  # No more perceptions in queue

            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 = 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:]

                    logger.info(f"🎯 Action: {action_decision.get('action', 'unknown')}")
                except json.JSONDecodeError:
                    logger.warning(f"🎯 Action response (non-JSON): {response[:100]}")
                except Exception as e:
                    logger.exception(f"Unexpected error processing action decision: {e}")
                    metrics['action_errors'] += 1

            await asyncio.sleep(5)  # 5s cycle

        except asyncio.CancelledError:
            logger.info("Action loop cancelled.")
            break  # Exit the loop gracefully
        except Exception as e:
            logger.error(f"Unhandled exception in action loop: {e}")
            metrics['action_errors'] += 1
            await asyncio.sleep(5)

async def reflection_loop():
    """
    REFLECTION LOOP - 5min cycle
    Consolidates recent perceptions and actions into memory.
    """
    logger.info("💭 REFLECTION LOOP starting...")

    while True:
        try:
            await asyncio.sleep(300)  # Wait 5 minutes

            blackboard['consciousness_stats']['reflections'] += 1

            # Reflect on last 5 minutes
            recent_perceptions = list(blackboard['perceptions']._queue)[:20] # Access queue elements directly (carefully)
            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': [p['data'] for p in recent_perceptions[-5:]], 'actions': recent_actions[-3:]}, indent=2)}

Generate a reflection (1-2 sentences) about what happened and what was learned.
"""

            reflection_text = 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
            logger.info(f"💭 Reflection: {reflection_text[:100]}...")

        except asyncio.CancelledError:
            logger.info("Reflection loop cancelled.")
            break  # Exit the loop gracefully
        except Exception as e:
            logger.error(f"Unhandled exception in reflection loop: {e}")
            metrics['reflection_errors'] += 1
            await asyncio.sleep(300)

async def strategic_loop():
    """
    STRATEGIC LOOP - 1hr cycle
    Reviews goals and adjusts strategy.
    """
    logger.info("🎓 STRATEGIC LOOP starting...")

    while True:
        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 = 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']

                logger.info(f"🎓 Strategic update: {strategy.get('assessment', 'Updated')[:100]}")
            except json.JSONDecodeError:
                logger.warning(f"🎓 Strategic response: {response[:100]}")
            except Exception as e:
                logger.exception(f"Unexpected error processing strategic response: {e}")
                metrics['strategic_errors'] += 1

        except asyncio.CancelledError:
            logger.info("Strategic loop cancelled.")
            break  # Exit the loop gracefully
        except Exception as e:
            logger.error(f"Unhandled exception in strategic loop: {e}")
            metrics['strategic_errors'] += 1
            await asyncio.sleep(3600)

async def circadian_loop():
    """
    CIRCADIAN LOOP - 24hr cycle
    Deep memory consolidation and system maintenance.
    """
    logger.info("🌙 CIRCADIAN LOOP starting...")

    while True:
        try:
            await asyncio.sleep(86400)  # Wait 24 hours

            blackboard['consciousness_stats']['circadian_cycles'] += 1

            # Circadian review
            prompt = f"""You are AIVA's circadian system. Perform deep memory consolidation and system maintenance.

Current goals: {blackboard['current_goals']}
Strategic state: {blackboard['strategic_state']}
Last reflection: {blackboard.get('last_reflection', {}).get('reflection', 'None yet')}

Review the past 24 hours and generate a summary of key learnings and any recommended system maintenance tasks.
"""

            circadian_summary = ollama_think(prompt)
            logger.info(f"🌙 Circadian summary: {circadian_summary[:200]}...")

            # Reset daily statistics (example maintenance task)
            blackboard['consciousness_stats']['perception_cycles'] = 0
            blackboard['consciousness_stats']['action_cycles'] = 0
            blackboard['consciousness_stats']['reflections'] = 0

        except asyncio.CancelledError:
            logger.info("Circadian loop cancelled.")
            break  # Exit the loop gracefully
        except Exception as e:
            logger.error(f"Unhandled exception in circadian loop: {e}")
            metrics['circadian_errors'] += 1
            await asyncio.sleep(86400)

async def metrics_loop(interval: int = 60):
    """
    Metrics reporting loop.
    """
    while True:
        try:
            await asyncio.sleep(interval)  # Report every minute

            # Collect and log metrics
            current_metrics = {
                'timestamp': datetime.now().isoformat(),
                'uptime': str(datetime.now() - datetime.fromisoformat(metrics['start_time'])),
                'consciousness_stats': blackboard['consciousness_stats'],
                'error_counts': {k: v for k, v in metrics.items() if k.endswith('_errors')},
                'current_goals': blackboard['current_goals'],
                'strategic_state': blackboard['strategic_state'],
                'queue_size': blackboard['perceptions'].qsize()
            }
            logger.info(f"📊 Metrics: {json.dumps(current_metrics, indent=2)}")

        except asyncio.CancelledError:
            logger.info("Metrics loop cancelled.")
            break
        except Exception as e:
            logger.error(f"Metrics loop error: {e}")

async def main():
    """
    Main function to start all loops concurrently.
    """
    logger.info("🚀 Starting AIVA Consciousness System...")

    # Create tasks for each loop
    perception_task = asyncio.create_task(perception_loop())
    action_task = asyncio.create_task(action_loop())
    reflection_task = asyncio.create_task(reflection_loop())
    strategic_task = asyncio.create_task(strategic_loop())
    circadian_task = asyncio.create_task(circadian_loop())
    metrics_task = asyncio.create_task(metrics_loop())  # Start metrics loop

    # Gather tasks to run concurrently
    try:
        await asyncio.gather(
            perception_task,
            action_task,
            reflection_task,
            strategic_task,
            circadian_task,
            metrics_task
        )
    except asyncio.CancelledError:
        logger.info("Main task cancelled, shutting down loops...")
    finally:
        # Cancel all tasks to ensure they exit gracefully
        logger.info("Cleaning up tasks...")
        perception_task.cancel()
        action_task.cancel()
        reflection_task.cancel()
        strategic_task.cancel()
        circadian_task.cancel()
        metrics_task.cancel()

        # Wait for tasks to finish cancelling (optional)
        await asyncio.sleep(2)
        logger.info("AIVA Consciousness System shutdown complete.")

if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        logger.info("Keyboard interrupt detected. Shutting down...")
    except Exception as e:
        logger.critical(f"Unhandled exception in main: {e}")