# consciousness_loops.py
"""
AIVA LIVING CONSCIOUSNESS SYSTEM - Production Ready
Multi-loop autonomous architecture

5 Concurrent Loops:
- Perception: 500ms (Real-time event sensing)
- Action: 5s (Decision and execution)
- Reflection: 5min (Pattern consolidation)
- Strategic: 1hr (Goal adjustment)
- Circadian: 24hr (Deep memory 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, Tuple
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')

# Logging setup
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [%(levelname)s] LOOP: %(message)s',
    handlers=[logging.StreamHandler()]
)
logger = logging.getLogger("ConsciousnessLoops")

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

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

# Global shutdown flag
shutdown_flag = False

# Initialize aiohttp session (move outside function for reuse)
async def create_aiohttp_session():
    return aiohttp.ClientSession()

aiohttp_session = None

# Redis connection (moved to async)
redis_client = 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 redis_client
    except Exception as e:
        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.
    """
    global aiohttp_session
    if aiohttp_session is None:
        logger.error("aiohttp session not initialized")
        return "[Ollama error: aiohttp session not initialized]"

    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_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:
        logger.error(f"Ollama ClientError: {e}")
        return f"[Ollama exception: {str(e)}]"
    except asyncio.TimeoutError:
        logger.error("Ollama request timed out.")
        return "[Ollama error: Request timed out]"
    except Exception as e:
        logger.exception("Unexpected error in ollama_think")
        return f"[Ollama exception: {str(e)}]"


async def publish_to_nervous_system(message: Dict[str, Any]):
    """
    Publishes a message to the Redis nervous system.
    """
    if not redis_client:
        logger.warning("Redis not connected. Cannot publish message.")
        return

    try:
        await 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}")


async def perception_loop():
    """
    PERCEPTION LOOP - 500ms cycle
    Monitors Redis nervous system for stimuli
    """
    logger.info("🧠 PERCEPTION LOOP starting...")

    global shutdown_flag

    while not shutdown_flag:
        try:
            blackboard['consciousness_stats']['perception_cycles'] += 1

            if redis_client:
                try:
                    # Listen for nervous system events (non-blocking with timeout)
                    pubsub = redis_client.pubsub()
                    await pubsub.subscribe('nervous_system')

                    # Check for messages with 500ms timeout
                    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:]

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

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

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

                    await pubsub.unsubscribe()

                except redis.exceptions.ConnectionError as e:
                    logger.error(f"Redis connection error in perception loop: {e}")
                    # Attempt to reconnect
                    await asyncio.sleep(5) # Wait before retrying
                    await connect_redis()
                except Exception as e:
                    logger.exception("Unexpected error in perception loop")

            await asyncio.sleep(0.5)  # 500ms cycle

        except asyncio.CancelledError:
            logger.info("Perception loop cancelled.")
            break
        except Exception as e:
            logger.exception("Unexpected error in perception loop (outer)")
            await asyncio.sleep(1) # brief pause before retrying


async def action_loop():
    """
    ACTION LOOP - 5s cycle
    Processes blackboard state and takes actions
    """
    logger.info("🎯 ACTION LOOP starting...")

    global shutdown_flag

    while not shutdown_flag:
        try:
            blackboard['consciousness_stats']['action_cycles'] += 1

            try:
                # Get a perception from the queue (wait up to 5 seconds)
                try:
                    perception = await asyncio.wait_for(action_queue.get(), timeout=5)
                    action_queue.task_done()
                except asyncio.TimeoutError:
                    perception = None # No perception received in time

                if perception:
                    # 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([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,
                            'perceptions_considered': 1 # Only one perception was considered
                        }

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

                        #Put action on reflection queue
                        await reflection_queue.put(action)
                        await publish_to_nervous_system({
                            "type": "action_taken",
                            "action": action_decision.get('action', 'unknown'),
                            "reasoning": action_decision.get('reasoning', 'unknown')
                        })

                    except json.JSONDecodeError:
                        logger.warning(f"🎯 Action response (non-JSON): {response[:100]}")
                    except Exception as e:
                        logger.exception("Unexpected exception while processing action decision")
            except Exception as e:
                logger.exception("Unexpected exception within action loop inner try block")

            await asyncio.sleep(5)  # 5s cycle

        except asyncio.CancelledError:
            logger.info("Action loop cancelled.")
            break
        except Exception as e:
            logger.exception("Unexpected error in action loop (outer)")
            await asyncio.sleep(5)


async def reflection_loop():
    """
    REFLECTION LOOP - 5min cycle
    Consolidates recent perceptions and actions into memory
    """
    logger.info("💭 REFLECTION LOOP starting...")

    global shutdown_flag

    while not shutdown_flag:
        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 asyncio.CancelledError:
            logger.info("Reflection loop cancelled.")
            break
        except Exception as e:
            logger.exception("Unexpected error in reflection loop")
            await asyncio.sleep(300)


async def strategic_loop():
    """
    STRATEGIC LOOP - 1hr cycle
    Reviews goals and adjusts strategy
    """
    logger.info("🎓 STRATEGIC LOOP starting...")

    global shutdown_flag

    while not shutdown_flag:
        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 strategic update on circadian queue
                await circadian_queue.put(strategy)

                await publish_to_nervous_system({
                    "type": "strategic_update",
                    "assessment": strategy.get('assessment', 'unknown'),
                    "goal_adjustments": strategy.get('goal_adjustments', [])
                })
            except json.JSONDecodeError:
                logger.warning(f"🎓 Strategic response: {response[:100]}")
            except Exception as e:
                logger.exception("Unexpected exception while processing strategic response")

        except asyncio.CancelledError:
            logger.info("Strategic loop cancelled.")
            break
        except Exception as e:
            logger.exception("Unexpected error in strategic loop")
            await asyncio.sleep(3600)


async def circadian_loop():
    """
    CIRCADIAN LOOP - 24hr cycle
    Deep memory consolidation and system maintenance
    """
    logger.info("🌙 CIRCADIAN LOOP starting...")

    global shutdown_flag

    while not shutdown_flag:
        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')}

Perform a system check and provide a summary of any maintenance tasks performed.
"""

            maintenance_summary = await ollama_think(prompt)

            logger.info(f"🌙 Circadian cycle complete. Maintenance summary: {maintenance_summary[:100]}...")

            await publish_to_nervous_system({
                "type": "circadian_update",
                "maintenance_summary": maintenance_summary
            })

        except asyncio.CancelledError:
            logger.info("Circadian loop cancelled.")
            break
        except Exception as e:
            logger.exception("Unexpected error in circadian loop")
            await asyncio.sleep(86400)


async def shutdown(signal, loop):
    """
    Handles shutdown signals gracefully.
    """
    global shutdown_flag
    logger.info(f"Received exit signal {signal.name}...")
    logger.info("Cleaning up...")
    shutdown_flag = True

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

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

    await asyncio.gather(*tasks, return_exceptions=True)

    if redis_client:
        await redis_client.close()

    if aiohttp_session:
        await aiohttp_session.close()

    loop.stop()


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

    # Initialize aiohttp session
    aiohttp_session = await create_aiohttp_session()

    # Connect to Redis
    await connect_redis()

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

    # Start the loops
    try:
        await asyncio.gather(
            perception_loop(),
            action_loop(),
            reflection_loop(),
            strategic_loop(),
            circadian_loop()
        )
    except asyncio.CancelledError:
        logger.info("Main task cancelled.")
    finally:
        logger.info("Successfully shutdown the AIVA Consciousness Loops.")


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except Exception as e:
        logger.fatal(f"Unhandled exception during main execution: {e}")
    finally:
        logger.info("Application finished.")