# consciousness_loops.py
"""
AIVA LIVING CONSCIOUSNESS SYSTEM - PRODUCTION READY
Multi-loop autonomous architecture with robust error handling and metrics.

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
import signal
from datetime import datetime, timedelta
from typing import Dict, Any, List, Optional
import redis.asyncio as aioredis
import aiohttp

# Configuration (Environment variables with defaults)
OLLAMA_HOST = os.getenv('OLLAMA_HOST', 'localhost')
OLLAMA_PORT = int(os.getenv('OLLAMA_PORT', '11434'))
AIVA_MODEL = os.getenv('AIVA_MODEL', 'qwen-long')
REDIS_URL = os.getenv('REDIS_URL', 'redis://localhost:6379')
LOG_LEVEL = os.getenv('LOG_LEVEL', 'INFO').upper()  # e.g., DEBUG, INFO, WARNING

# Logging setup
logging.basicConfig(
    level=LOG_LEVEL,
    format='%(asctime)s [%(levelname)s] CONSCIOUSNESS: %(message)s',
    handlers=[logging.StreamHandler()]
)
logger = logging.getLogger(__name__)

# Global shutdown event
shutdown_event = asyncio.Event()

# Shared blackboard for inter-loop communication (using asyncio.Queue)
class Blackboard:
    def __init__(self):
        self.perceptions: asyncio.Queue = asyncio.Queue(maxsize=100)  # Bounded queue
        self.actions_taken: asyncio.Queue = asyncio.Queue(maxsize=50) # Bounded queue
        self.current_goals: List[str] = ['Understand Genesis', 'Process patent knowledge', 'Coordinate MVP']
        self.strategic_state: str = 'initializing'
        self.last_reflection: Optional[Dict[str, Any]] = None
        self.consciousness_stats: Dict[str, int] = {
            'perception_cycles': 0,
            'action_cycles': 0,
            'reflections': 0,
            'strategic_reviews': 0,
            'circadian_cycles': 0
        }

blackboard = Blackboard() # Instantiate the Blackboard

# Metrics tracking
metrics = {
    'perception_errors': 0,
    'action_errors': 0,
    'reflection_errors': 0,
    'strategic_errors': 0,
    'circadian_errors': 0,
    'total_cycles': 0
}

# Redis client (using aioredis)
redis_client: Optional[aioredis.Redis] = None


async def connect_redis():
    global redis_client
    try:
        redis_client = aioredis.from_url(REDIS_URL)
        await redis_client.ping()
        logger.info(f"✅ Connected to Redis: {REDIS_URL}")
        return True
    except Exception as e:
        logger.error(f"⚠️ Redis connection failed: {e}")
        redis_client = None
        return False


async def ollama_think(prompt: str, model: str = AIVA_MODEL) -> str:
    """
    Asynchronous call to 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
        }
    }

    async with aiohttp.ClientSession() as session:
        try:
            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 aiohttp.ClientError as e:
            return f"[Ollama connection error: {str(e)}]"
        except asyncio.TimeoutError:
            return "[Ollama timeout error]"
        except Exception as e:
            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 not shutdown_event.is_set():
        try:
            blackboard.consciousness_stats['perception_cycles'] += 1
            metrics['total_cycles'] += 1

            if redis_client:
                try:
                    async with redis_client.pubsub() as pubsub:
                        await pubsub.subscribe('genesis:nervous_system')

                        try:
                            message = await pubsub.get_message(timeout=0.5, ignore_subscribe_messages=True)

                            if message and message['type'] == 'message':
                                try:
                                    data = json.loads(message['data'].decode('utf-8'))

                                    perception = {
                                        'timestamp': datetime.now().isoformat(),
                                        'type': 'redis_event',
                                        'data': data
                                    }

                                    try:
                                        await blackboard.perceptions.put(perception)
                                    except asyncio.QueueFull:
                                        logger.warning("Perception queue is full. Dropping oldest perception.")
                                        blackboard.perceptions._queue.popleft() # Remove oldest item
                                        await blackboard.perceptions.put(perception)


                                    logger.info(f"👁️ Perceived: {data.get('type', 'unknown')}")

                                except json.JSONDecodeError as e:
                                    logger.warning(f"Invalid JSON received from Redis: {e}")

                        except asyncio.TimeoutError:
                            pass  # No message received within the timeout

                except aioredis.exceptions.RedisError as e:
                    logger.error(f"Redis error in perception loop: {e}")
                    metrics['perception_errors'] += 1
                    await asyncio.sleep(5) # Wait before retrying Redis

            await asyncio.sleep(0.5)  # 500ms cycle

        except Exception as e:
            logger.exception(f"⚠️ Perception error: {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 not shutdown_event.is_set():
        try:
            blackboard.consciousness_stats['action_cycles'] += 1
            metrics['total_cycles'] += 1

            # Check for new perceptions to act on (using 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 that the queue item is processed
                except asyncio.QueueEmpty:
                    break  # Queue is empty

            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)
                    }
                    try:
                        await blackboard.actions_taken.put(action)
                    except asyncio.QueueFull:
                        logger.warning("Action queue is full. Dropping oldest action.")
                        blackboard.actions_taken._queue.popleft()
                        await blackboard.actions_taken.put(action)


                    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"🎯 Error processing action decision: {e}")

            await asyncio.sleep(5)  # 5s cycle

        except Exception as e:
            logger.exception(f"⚠️ Action error: {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 not shutdown_event.is_set():
        try:
            await asyncio.sleep(300)  # Wait 5 minutes

            blackboard.consciousness_stats['reflections'] += 1
            metrics['total_cycles'] += 1

            # Reflect on last 5 minutes (consume the queues)
            recent_perceptions: List[Dict[str, Any]] = []
            while not blackboard.perceptions.empty():
                try:
                    recent_perceptions.append(await blackboard.perceptions.get_nowait())
                    blackboard.perceptions.task_done()
                except asyncio.QueueEmpty:
                    break

            recent_actions: List[Dict[str, Any]] = []
            while not blackboard.actions_taken.empty():
                try:
                    recent_actions.append(await blackboard.actions_taken.get_nowait())
                    blackboard.actions_taken.task_done()
                except asyncio.QueueEmpty:
                    break

            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
            logger.info(f"💭 Reflection: {reflection_text[:100]}...")

        except Exception as e:
            logger.exception(f"⚠️ Reflection error: {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 not shutdown_event.is_set():
        try:
            await asyncio.sleep(3600)  # Wait 1 hour

            blackboard.consciousness_stats['strategic_reviews'] += 1
            metrics['total_cycles'] += 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.last_reflection.get('reflection', 'None yet') if blackboard.last_reflection else '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']

                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"🎓 Error processing strategic response: {e}")

        except Exception as e:
            logger.exception(f"⚠️ Strategic error: {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 not shutdown_event.is_set():
        try:
            await asyncio.sleep(86400)  # Wait 24 hours

            blackboard.consciousness_stats['circadian_cycles'] += 1
            metrics['total_cycles'] += 1

            # Deep memory consolidation and system maintenance
            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.last_reflection.get('reflection', 'None yet') if blackboard.last_reflection else 'None yet'}
Consciousness stats: {blackboard.consciousness_stats}

Generate a summary of key learnings, suggest system improvements, and propose long-term goals (respond in JSON):
{{
  "summary": "<summary>",
  "improvements": ["<improvement1>", "<improvement2>", ...],
  "long_term_goals": ["<goal1>", "<goal2>", ...]
}}
"""

            response = await ollama_think(prompt)

            try:
                circadian_data = json.loads(response)
                logger.info(f"🌙 Circadian update: {circadian_data.get('summary', 'Completed')[:100]}")

                # Example: Update goals based on long-term goals
                if 'long_term_goals' in circadian_data:
                    blackboard.current_goals = circadian_data['long_term_goals']

            except json.JSONDecodeError:
                logger.warning(f"🌙 Circadian response (non-JSON): {response[:100]}")
            except Exception as e:
                logger.exception(f"🌙 Error processing circadian data: {e}")

        except Exception as e:
            logger.exception(f"⚠️ Circadian error: {e}")
            metrics['circadian_errors'] += 1
            await asyncio.sleep(86400)


async def metrics_loop():
    """Periodically log system metrics."""
    while not shutdown_event.is_set():
        logger.info(f"System Metrics: {metrics} | Consciousness Stats: {blackboard.consciousness_stats}")
        await asyncio.sleep(600)  # Log every 10 minutes


async def shutdown(signal, loop):
    """Handles shutdown signals gracefully."""
    logger.info(f"Received exit signal {signal.name}...")
    logger.info("Cleaning up...")
    shutdown_event.set()  # Signal all loops to stop

    tasks = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()]
    for task in tasks:
        task.cancel()

    logger.info(f"Cancelling {len(tasks)} outstanding tasks")
    await asyncio.gather(*tasks, return_exceptions=True)
    logger.info(f"Flushing Redis connection...")
    if redis_client:
        await redis_client.close()
    loop.stop()


async def main():
    """Main function to start the consciousness loops."""
    loop = asyncio.get_running_loop()

    # Handle shutdown signals
    signals = (signal.SIGHUP, signal.SIGTERM, signal.SIGINT)
    for s in signals:
        loop.add_signal_handler(
            s, lambda s=s: asyncio.create_task(shutdown(s, loop)))

    # Connect to Redis
    redis_connected = await connect_redis()
    if not redis_connected:
        logger.error("Failed to connect to Redis. Exiting.")
        return

    # Start the loops
    try:
        await asyncio.gather(
            perception_loop(),
            action_loop(),
            reflection_loop(),
            strategic_loop(),
            circadian_loop(),
            metrics_loop() # Monitoring loop
        )
    except asyncio.CancelledError:
        logger.info("Asyncio gather cancelled")
    finally:
        logger.info("Consciousness loops finished.")


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        logger.info("Keyboard interrupt received. Shutting down.")
    finally:
        logger.info("Application finished.")