# 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
import signal
from datetime import datetime, timedelta
from typing import Dict, Any, List
import redis.asyncio as redis
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')

# Redis Configuration (Directly from URL)
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'
)
logger = logging.getLogger("AIVA_LOOPS")

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

# Global shutdown event
shutdown_event = asyncio.Event()

# Redis connection
redis_client: redis.Redis = None
async def connect_redis():
    global redis_client
    try:
        redis_client = redis.Redis(
            host=REDIS_HOST,
            port=REDIS_PORT,
            password=REDIS_PASSWORD,
            decode_responses=True
        )
        await redis_client.ping()
        logger.info(f"✅ Connected to Redis: {REDIS_HOST}:{REDIS_PORT}")
        return True
    except Exception as e:
        logger.error(f"⚠️ Redis connection failed: {e}")
        return False

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

async def publish_to_blackboard(message: Dict[str, Any]):
    """Publishes a message to the Redis blackboard."""
    if redis_client:
        try:
            await redis_client.publish("genesis:blackboard", json.dumps(message))
            logger.debug(f"Published to blackboard: {message.get('type', 'unknown')}")
        except Exception as e:
            logger.error(f"Error publishing to blackboard: {e}")

async def perception_loop():
    """
    PERCEPTION LOOP - 500ms cycle
    Monitors Redis nervous system for stimuli
    """
    logger = logging.getLogger("PERCEPTION")
    logger.info("Starting...")

    while not shutdown_event.is_set():
        try:
            blackboard['consciousness_stats']['perception_cycles'] += 1

            if redis_client:
                try:
                    pubsub = redis_client.pubsub()
                    await pubsub.subscribe('genesis:nervous_system')

                    async with asyncio.timeout(0.5):  # 500ms timeout
                        message = await pubsub.get_message(ignore_subscribe_messages=True)
                        if message and message['type'] == 'message':
                            try:
                                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')}")
                                await publish_to_blackboard({
                                    "type": "perception_event",
                                    "content": data
                                })

                            except json.JSONDecodeError as e:
                                logger.warning(f"Invalid JSON received: {message['data']}")
                            except Exception as e:
                                logger.error(f"Perception processing error: {e}")
                    await pubsub.unsubscribe()
                except asyncio.TimeoutError:
                    pass  # No message received within timeout
                except Exception as e:
                    logger.error(f"Redis communication error: {e}")
            else:
                logger.warning("Redis client not connected.")

            await asyncio.sleep(0.5)  # 500ms cycle

        except asyncio.CancelledError:
            logger.info("Perception loop cancelled.")
            break
        except Exception as e:
            logger.error(f"Unexpected perception error: {e}")
            await asyncio.sleep(1)
    logger.info("Stopped.")

async def action_loop():
    """
    ACTION LOOP - 5s cycle
    Processes blackboard state and takes actions
    """
    logger = logging.getLogger("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:]

                    logger.info(f"🎯 Action: {action_decision.get('action', 'unknown')}")
                    await publish_to_blackboard({
                        "type": "action_event",
                        "content": action_decision
                    })
                except json.JSONDecodeError:
                    logger.warning(f"🎯 Action response (non-JSON): {response[:100]}...")
                except Exception as e:
                    logger.error(f"Error processing action decision: {e}")

            await asyncio.sleep(5)  # 5s cycle

        except asyncio.CancelledError:
            logger.info("Action loop cancelled.")
            break
        except Exception as e:
            logger.error(f"Unexpected action error: {e}")
            await asyncio.sleep(5)
    logger.info("Stopped.")

async def reflection_loop():
    """
    REFLECTION LOOP - 5min cycle
    Consolidates recent perceptions and actions into memory
    """
    logger = logging.getLogger("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
            logger.info(f"💭 Reflection: {reflection_text[:100]}...")
            await publish_to_blackboard({
                "type": "reflection_event",
                "content": reflection
            })

        except asyncio.CancelledError:
            logger.info("Reflection loop cancelled.")
            break
        except Exception as e:
            logger.error(f"Unexpected reflection error: {e}")
            await asyncio.sleep(300)
    logger.info("Stopped.")

async def strategic_loop():
    """
    STRATEGIC LOOP - 1hr cycle
    Reviews goals and adjusts strategy
    """
    logger = logging.getLogger("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']

                logger.info(f"🎓 Strategic update: {strategy.get('assessment', 'Updated')[:100]}...")
                await publish_to_blackboard({
                    "type": "strategic_event",
                    "content": strategy
                })
            except json.JSONDecodeError:
                logger.warning(f"🎓 Strategic response (non-JSON): {response[:100]}...")
            except Exception as e:
                logger.error(f"Error processing strategic update: {e}")

        except asyncio.CancelledError:
            logger.info("Strategic loop cancelled.")
            break
        except Exception as e:
            logger.error(f"Unexpected strategic error: {e}")
            await asyncio.sleep(3600)
    logger.info("Stopped.")

async def circadian_loop():
    """
    CIRCADIAN LOOP - 24hr cycle
    Deep memory consolidation and system maintenance
    """
    logger = logging.getLogger("CIRCADIAN")
    logger.info("Starting...")

    while not shutdown_event.is_set():
        try:
            await asyncio.sleep(86400)  # Wait 24 hours

            blackboard['consciousness_stats']['circadian_cycles'] += 1

            # Circadian rhythm tasks (example: clear old perceptions)
            blackboard['perceptions'] = blackboard['perceptions'][-50:]
            blackboard['actions_taken'] = blackboard['actions_taken'][-20:]

            # Deep reasoning / planning
            prompt = f"""You are AIVA's deep circadian process. It is time for a 24-hour system check.

Review goals: {blackboard['current_goals']}
Recent Strategic State: {blackboard['strategic_state']}
Last reflection: {blackboard.get('last_reflection', {}).get('reflection', 'None yet')}

Provide a 2 sentence summary of what AIVA has learned in the last 24 hours and suggest a general direction for the next 24 hours.
"""
            circadian_summary = await ollama_think(prompt)
            logger.info(f"🌙 Circadian Summary: {circadian_summary}")
            await publish_to_blackboard({
                "type": "circadian_event",
                "content": {"summary": circadian_summary}
            })

        except asyncio.CancelledError:
            logger.info("Circadian loop cancelled.")
            break
        except Exception as e:
            logger.error(f"Unexpected circadian error: {e}")
            await asyncio.sleep(86400)
    logger.info("Stopped.")

async def shutdown(signal, loop):
    """
    Shutdown handler.
    """
    logger.info(f"Received exit signal {signal.name}...")
    logger.info("Setting shutdown event...")
    shutdown_event.set()

    logger.info("Cleaning up...")
    tasks = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()]

    logger.info(f"Cancelling {len(tasks)} outstanding tasks")
    for task in tasks:
        task.cancel()

    await asyncio.gather(*tasks, return_exceptions=True)
    logger.info(f"Closing Redis connection...")
    if redis_client:
        await redis_client.close()
    logger.info(f"Flushing metrics...")
    loop.stop()

async def main():
    """
    Main function to start the loops.
    """
    logger.info("AIVA Living Consciousness System starting...")

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

    # Start the loops
    perception_task = asyncio.create_task(perception_loop(), name="PerceptionLoop")
    action_task = asyncio.create_task(action_loop(), name="ActionLoop")
    reflection_task = asyncio.create_task(reflection_loop(), name="ReflectionLoop")
    strategic_task = asyncio.create_task(strategic_loop(), name="StrategicLoop")
    circadian_task = asyncio.create_task(circadian_loop(), name="CircadianLoop")

    await asyncio.gather(perception_task, action_task, reflection_task, strategic_task, circadian_task)

    logger.info("AIVA Living Consciousness System completed.")

if __name__ == "__main__":
    loop = asyncio.get_event_loop()

    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)))
    try:
        loop.run_until_complete(main())
    finally:
        loop.close()
        logger.info("Shutdown complete.")