# consciousness_loops.py
import asyncio
import time
import json
import os
from datetime import datetime, timedelta
from typing import Dict, Any, List, Tuple
import redis.asyncio as aioredis
import aiohttp
import logging

# Configuration (use environment variables)
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")

# Shared Queues for inter-loop communication
perception_queue = asyncio.Queue()
action_queue = asyncio.Queue()
reflection_queue = asyncio.Queue()
strategic_queue = asyncio.Queue()

# Global State (avoid global state as much as possible, but necessary for shutdown)
running = True

async def shutdown():
    """
    Graceful shutdown handler.
    """
    global running
    running = False
    logger.info("Shutting down loops...")

async def ollama_think(prompt: str, model: str = AIVA_MODEL) -> str:
    """
    Asynchronous call to Ollama for thinking/reasoning
    """
    url = f"http://{OLLAMA_HOST}:{OLLAMA_PORT}/api/generate"
    payload = {
        "model": model,
        "prompt": prompt,
        "stream": False,
        "options": {
            "temperature": 0.7,
            "num_predict": 500
        }
    }

    try:
        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 aiohttp.ClientError as e:
        return f"[Ollama connection error: {str(e)}]"
    except Exception as e:
        return f"[Ollama exception: {str(e)}]"

async def perception_loop(redis_client: aioredis.Redis, blackboard: Dict[str, Any]):
    """
    PERCEPTION LOOP - 500ms cycle
    Monitors Redis nervous system for stimuli
    """
    logger.info("🧠 PERCEPTION LOOP starting...")

    perception_cycles = 0
    while running:
        try:
            perception_cycles += 1
            blackboard['consciousness_stats']['perception_cycles'] = perception_cycles

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

                try:
                    message = await 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')}")
                        await perception_queue.put(perception)  # Send to action loop
                finally:
                    await pubsub.unsubscribe()
                    await pubsub.close()

            except aioredis.exceptions.RedisError as e:
                logger.error(f"Redis error in perception loop: {e}")

            await asyncio.sleep(0.5)  # 500ms cycle

        except Exception as e:
            logger.exception(f"⚠️ Perception error: {e}")
            await asyncio.sleep(1)

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

    action_cycles = 0
    while running:
        try:
            action_cycles += 1
            blackboard['consciousness_stats']['action_cycles'] = action_cycles

            try:
                # Get recent perceptions from the queue
                recent_perceptions = []
                while not perception_queue.empty():
                    recent_perceptions.append(perception_queue.get_nowait())

                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 reflection_queue.put(action)  # Send to reflection loop

                    except json.JSONDecodeError:
                        logger.warning(f"🎯 Action response (non-JSON): {response[:100]}")
            except asyncio.QueueEmpty:
                pass # No perceptions to process

            await asyncio.sleep(5)  # 5s cycle

        except Exception as e:
            logger.exception(f"⚠️ Action error: {e}")
            await asyncio.sleep(5)

async def reflection_loop(blackboard: Dict[str, Any]):
    """
    REFLECTION LOOP - 5min cycle
    Consolidates recent perceptions and actions into memory
    """
    logger.info("💭 REFLECTION LOOP starting...")

    reflections = 0
    while running:
        try:
            await asyncio.sleep(300)  # Wait 5 minutes

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

            # Get recent actions from the queue
            recent_actions = []
            while not reflection_queue.empty():
                recent_actions.append(reflection_queue.get_nowait())

            # Reflect on last 5 minutes
            recent_perceptions = blackboard['perceptions'][-20:]

            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 strategic_queue.put(reflection)  # Send to strategic loop

        except Exception as e:
            logger.exception(f"⚠️ Reflection error: {e}")
            await asyncio.sleep(300)

async def strategic_loop(blackboard: Dict[str, Any]):
    """
    STRATEGIC LOOP - 1hr cycle
    Reviews goals and adjusts strategy
    """
    logger.info("🎓 STRATEGIC LOOP starting...")

    strategic_reviews = 0
    while running:
        try:
            await asyncio.sleep(3600)  # Wait 1 hour

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

            # 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]}")
            except json.JSONDecodeError:
                logger.warning(f"🎓 Strategic response: {response[:100]}")

        except Exception as e:
            logger.exception(f"⚠️ Strategic error: {e}")
            await asyncio.sleep(3600)

async def circadian_loop(blackboard: Dict[str, Any]):
    """
    CIRCADIAN LOOP - 24hr cycle
    Deep memory consolidation and system maintenance
    """
    logger.info("🌙 CIRCADIAN LOOP starting...")

    circadian_cycles = 0
    while running:
        try:
            await asyncio.sleep(86400)  # Wait 24 hours

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

            # Deep memory consolidation (example: summarize all reflections)
            all_reflections = [r.get('reflection', '') for r in blackboard.get('last_reflection', []) if isinstance(blackboard.get('last_reflection'), list)]
            summary_prompt = f"""
            You are AIVA's deep consolidation system.  Summarize the following reflections to improve long-term understanding:

            Reflections: {all_reflections}

            Provide a concise summary (max 3 sentences) of key insights and learning.
            """

            summary = await ollama_think(summary_prompt)
            logger.info(f"🌙 Circadian summary: {summary}")

            # System maintenance (example: clear old perceptions)
            blackboard['perceptions'] = blackboard['perceptions'][-50:]  # Keep only last 50
            logger.info("🌙 System maintenance complete: Old perceptions cleared.")

        except Exception as e:
            logger.exception(f"⚠️ Circadian error: {e}")
            await asyncio.sleep(86400)

async def main():
    """
    Main function to orchestrate the consciousness loops.
    """
    # Shared blackboard for inter-loop communication
    blackboard = {
        '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
        }
    }

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

    # Create tasks for each loop
    perception_task = asyncio.create_task(perception_loop(redis_client, blackboard))
    action_task = asyncio.create_task(action_loop(blackboard))
    reflection_task = asyncio.create_task(reflection_loop(blackboard))
    strategic_task = asyncio.create_task(strategic_loop(blackboard))
    circadian_task = asyncio.create_task(circadian_loop(blackboard))

    # Handle shutdown signals (Ctrl+C)
    try:
        await asyncio.gather(perception_task, action_task, reflection_task, strategic_task, circadian_task)
    except asyncio.CancelledError:
        logger.info("Tasks cancelled.")
    finally:
        await redis_client.close()
        logger.info("Redis connection closed.")
        logger.info("Consciousness loops stopped.")

if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        logger.info("Keyboard interrupt detected. Shutting down...")
        asyncio.run(shutdown()) # Shutdown all loops