# consciousness_loops.py
"""
AIVA LIVING CONSCIOUSNESS SYSTEM - Production Ready
Multi-loop autonomous architecture with async/await, graceful shutdown,
error recovery, metrics, and inter-loop communication.

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
import requests
import aiohttp

# 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 (using a more robust method)
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
LOG_PATH = "consciousness_loops.log"
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [%(levelname)s] CLOOPS: %(message)s',
    handlers=[
        logging.FileHandler(LOG_PATH),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger("CLOOPS")


# 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
action_queue: asyncio.Queue = asyncio.Queue()
reflection_queue: asyncio.Queue = asyncio.Queue()
strategic_queue: asyncio.Queue = asyncio.Queue()
circadian_queue: asyncio.Queue = asyncio.Queue()


# Redis connection (using aioredis for async)
redis_client: Optional[redis.Redis] = None

async def redis_connect():
    global redis_client
    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}")
        redis_client = None

async def ollama_think(prompt: str, model: str = AIVA_MODEL) -> str:
    """
    Call Ollama for thinking/reasoning using aiohttp.
    """
    try:
        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:
            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:
        logger.exception("Ollama think error")
        return f"[Ollama exception: {str(e)}]"


async def publish_to_nervous_system(message: Dict[str, Any]):
    """Publishes a message to the 'nervous_system' Redis channel."""
    if redis_client:
        try:
            await asyncio.to_thread(redis_client.publish, 'nervous_system', json.dumps(message))
            logger.debug(f"Published to nervous_system: {message.get('type', 'unknown')}")
        except Exception as e:
            logger.error(f"Error publishing to nervous system: {e}")
    else:
        logger.warning("Redis not connected, cannot publish to nervous_system.")


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
                        }

                        blackboard['perceptions'].append(perception)

                        # Keep only last 100 perceptions
                        if len(blackboard['perceptions']) > 100:
                            blackboard['perceptions'] = blackboard['perceptions'][-100:]

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

                        # Put perception on the action queue
                        await action_queue.put(perception)

                    pubsub.close()

                except Exception as e:
                    logger.exception("Perception loop Redis error")

            await asyncio.sleep(0.5)  # 500ms cycle

        except Exception as e:
            logger.exception("Perception loop error")
            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
            try:
                perception = await asyncio.wait_for(action_queue.get(), timeout=5.0)
                action_queue.task_done()
            except asyncio.TimeoutError:
                perception = None

            if perception:
                # Decide on action using Ollama
                prompt = f"""You are AIVA's action system. Based on a recent perception, decide what action to take.

Recent perception: {json.dumps(perception, 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,
                        'perception_considered': perception
                    }

                    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')}")

                    # Example action: Publish feedback to nervous system
                    if action_decision.get('action') == 'publish_feedback':
                        feedback_message = {
                            'type': 'aiva_action_feedback',
                            'content': action_decision.get('reasoning', 'Action taken.'),
                            'source': 'aiva_action_loop'
                        }
                        await publish_to_nervous_system(feedback_message)

                    # Put action on reflection queue
                    await reflection_queue.put(action)
                except json.JSONDecodeError:
                    logger.warning(f"🎯 Action response (non-JSON): {response[:100]}")
                except Exception as e:
                    logger.exception("Action loop processing error")
            else:
                logger.debug("🎯 No new perceptions to act on.")

            await asyncio.sleep(5)  # 5s cycle

        except Exception as e:
            logger.exception("Action loop error")
            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 = 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
            logger.info(f"💭 Reflection: {reflection_text[:100]}...")

            # Put reflection on strategic queue
            await strategic_queue.put(reflection)

        except Exception as e:
            logger.exception("Reflection loop error")
            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 = 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]}")

                # Put strategy on circadian queue
                await circadian_queue.put(strategy)
            except json.JSONDecodeError:
                logger.warning(f"🎓 Strategic response: {response[:100]}")
            except Exception as e:
                logger.exception("Strategic loop processing error")

        except Exception as e:
            logger.exception("Strategic loop error")
            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

            # Deep consolidation - simplified for example
            prompt = f"""You are AIVA's circadian system. Perform a deep consolidation of memory and system maintenance.

Last strategic review: {blackboard.get('strategic_state', 'None')}
Last reflection: {blackboard.get('last_reflection', {}).get('reflection', 'None yet')}

Summarize key learnings from the past 24 hours and suggest any long-term maintenance tasks.
"""

            summary = await ollama_think(prompt)
            logger.info(f"🌙 Circadian summary: {summary[:200]}...")

            # Example maintenance: Clear old perceptions
            blackboard['perceptions'] = blackboard['perceptions'][-20:]
            blackboard['actions_taken'] = blackboard['actions_taken'][-10:]
            logger.info("🌙 Perceptions and actions trimmed for maintenance.")

        except Exception as e:
            logger.exception("Circadian loop error")
            await asyncio.sleep(86400)


async def main():
    """
    Main function to start all loops concurrently.
    """
    await redis_connect()

    # 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())

    # Handle shutdown signals gracefully
    def signal_handler(sig, frame):
        logger.info(f"Received signal {sig}. Shutting down...")
        perception_task.cancel()
        action_task.cancel()
        reflection_task.cancel()
        strategic_task.cancel()
        circadian_task.cancel()

    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, 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
        )
    except asyncio.CancelledError:
        logger.info("Tasks cancelled.")
    finally:
        if redis_client:
            redis_client.close()
            logger.info("Redis connection closed.")
        logger.info("AIVA Consciousness Loops system shutdown.")


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        logger.info("Keyboard interrupt received. Shutting down...")
    except Exception as e:
        logger.exception("Unhandled exception in main")
    finally:
        logger.info("Exiting AIVA Consciousness Loops system.")