"""
Genesis Swarm Manager
Orchestrates low-cost 24/7 Gemini Flash agents to evolve the system.
"""

import sys
import os
import time
from typing import List, Dict, Any
from datetime import datetime

# Add root to path for imports
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from swarms.config import swarm_config
import google.generativeai as genai

class GenesisSwarm:
    def __init__(self, daily_budget_usd: float = 10.0):
        self.coordinator_model = swarm_config.get_model("reasoning") # Flash Thinking
        self.worker_model = swarm_config.get_model("fast") # Flash 2.0
        self.creative_model = swarm_config.get_model("creative") # Exp 1206
        
        self.daily_budget = daily_budget_usd
        self.current_spend = 0.0
        self.token_usage = {"input": 0, "output": 0}
        
        # Approximate Costs (Flash 2.0)
        self.COST_INPUT_1M = 0.10
        self.COST_OUTPUT_1M = 0.40
        
        # Advanced Memory Hardwiring (Qdrant)
        self.qdrant_client = None
        if swarm_config.qdrant:
            try:
                from qdrant_client import QdrantClient
                from fastembed import TextEmbedding
                self.qdrant_client = QdrantClient(
                    url=swarm_config.qdrant['url'],
                    api_key=swarm_config.qdrant['api_key']
                )
                from swarms.master_context_indexer import EMBEDDING_MODEL # Reuse model name
                self.embedding_model = TextEmbedding(model_name=EMBEDDING_MODEL)
                print("[Memory] Advanced Memory System Connected.")
            except Exception as e:
                print(f"[Memory] Failed to connect to Qdrant: {e}")

        print(f"[Swarm] Initialized. Budget: ${self.daily_budget}/day")

    def retrieve_context(self, query: str, limit: int = 5) -> str:
        """Retrieves system context from Qdrant RAG."""
        if not self.qdrant_client:
            return ""
        
        try:
            vector = list(self.embedding_model.embed([query]))[0].tolist()
            results = self.qdrant_client.search(
                collection_name="genesis_master_context",
                query_vector=vector,
                limit=limit
            )
            
            context_blocks = []
            for res in results:
                path = res.payload.get("path", "unknown")
                text = res.payload.get("text", "")
                context_blocks.append(f"--- SOURCE: {path} ---\n{text}")
            
            return "\n\n".join(context_blocks)
        except Exception as e:
            print(f"⚠️ [Memory] Retrieval failed: {e}")
            return ""

    def track_cost(self, usage_metadata):
        """Update system cost based on token usage."""
        if not usage_metadata:
            return

        in_tokens = usage_metadata.prompt_token_count
        out_tokens = usage_metadata.candidates_token_count
        
        self.token_usage["input"] += in_tokens
        self.token_usage["output"] += out_tokens
        
        cost = (in_tokens / 1_000_000 * self.COST_INPUT_1M) + \
               (out_tokens / 1_000_000 * self.COST_OUTPUT_1M)
               
        self.current_spend += cost
        
        if self.current_spend >= self.daily_budget:
            print(f"⚠️ [Swarm] BUDGET HIT: ${self.current_spend:.4f} / ${self.daily_budget}")
            raise Exception("Daily Budget Exceeded")

    def spawn_agent(self, role: str, task: str, use_memory: bool = True) -> str:
        """
        Spawns a specialized Gemini Flash agent to perform a task.
        """
        print(f"  > [Swarm] Spawning '{role}' agent for: {task[:50]}...")
        
        context = ""
        if use_memory:
            context = self.retrieve_context(task)
            if context:
                print(f"    🧠 [Memory] Retrieved context from system knowledge.")

        system_instruction = f"""
        You are a specialized component of the CLAIM (Continuous Learning AI Swarm).
        Role: {role}
        Objective: Execute the task autonomously and return a structured result.
        Context: You are running on the Genesis System (Windows).
        
        RELEVANT SYSTEM KNOWLEDGE:
        {context}
        """
        
        try:
            # Select model based on role complexity
            if role in ["Architect", "Planner", "Lead Architect", "CEO"]:
                model = self.coordinator_model
                model_name = "thinking"
            elif role in ["Writer", "Strategist"]:
                model = self.creative_model
                model_name = "creative"
            else:
                model = self.worker_model
                model_name = "flash"
            
            start_time = time.time()
            response = model.generate_content(
                f"{system_instruction}\n\nTask: {task}"
            )
            elapsed = time.time() - start_time
            
            # Track Usage
            self.track_cost(response.usage_metadata)
            
            print(f"    - [{model_name}] Finished in {elapsed:.2f}s | Cost so far: ${self.current_spend:.4f}")
            return response.text
            
        except Exception as e:
            return f"Agent Failure: {e}"

    def _get_scout_prompt(self) -> str:
        """Generate the scout prompt for evolution cycle."""
        return """
        Scan the Genesis System for evolution opportunities:

        1. CHECK TASK QUEUE: Review loop/tasks.json for pending RWL stories
        2. KNOWLEDGE GAPS: Identify missing entities in KNOWLEDGE_GRAPH/
        3. CODE HEALTH: Look for TODOs, FIXMEs, or deprecated patterns
        4. INTEGRATION STATUS: Verify Elestio services (Postgres, Redis, Qdrant) are healthy
        5. REVENUE PIPELINE: Check ReceptionistAI progress and blockers

        Return a prioritized list of 3 actionable items for this cycle.
        Format: JSON array with {priority, task, rationale}
        """

    def run_evolution_cycle(self):
        """
        Main heartbeat method.
        1. Scout (Find info)
        2. Analyze (Process info)
        3. Evolve (Propose changes)
        """
        print("\n🔄 [Cycle] Starting Evolution Loop...")
        
        # [HEARTBEAT] Update timestamp for Self-Healing Node
        try:
            with open("swarms/heartbeat.txt", "w") as f:
                f.write(str(time.time()))
        except Exception as e:
            print(f"⚠️ [Heartbeat] Failed to update: {e}")
        
        # 1. Scout
        scout_report = self.spawn_agent(
            "Lead Architect", 
            self._get_scout_prompt()
        )
        print(f"  [Report] {scout_report.strip()}")
        
        # 2. Architect
        plan = self.spawn_agent(
            "Architect",
            f"Based on this report: '{scout_report}', define a JSON plan for the next hour."
        )
        print(f"  [Plan] {plan.strip()}")

    def trigger_nervous_system(self, signal: str, payload: Dict[str, Any] = None):
        """
        Connects the Swarm Brain to the n8n Nervous System.
        Sends a signal to the Elestio n8n instance.
        """
        webhook_url = "https://n8n-genesis-u50607.vm.elestio.app/webhook"
        
        try:
            import requests
            data = {"signal": signal, "payload": payload or {}, "timestamp": datetime.now().isoformat()}
            # Send async to avoid blocking the brain (timeout=1s)
            response = requests.post(webhook_url, json=data, timeout=2)
            
            if response.status_code == 200:
                print(f"⚡ [Nervous System] Signal fired: {signal} (200 OK)")
            else:
                print(f"⚠️ [Nervous System] Signal failed: {signal} ({response.status_code})")
                
        except Exception as e:
            print(f"⚠️ [Nervous System] Connection Error: {e}")

    def run_high_velocity_cycle(self):
        """
        HIGH VELOCITY MODE ($10/day Target)
        Replaces linear loop with 'Boardroom & Factory' parallel agents.
        """
        print("\n🚀 [High Velocity] Initiating Boardroom & Factory Cycle...")
        
        # [HEARTBEAT]
        try:
            with open("swarms/heartbeat.txt", "w") as f:
                f.write(str(time.time()))
        except: pass

        # STAGE 1: THE BOARDROOM (3 Agents Debate)
        print("  🏛️ [Boardroom] Convening Strategy Council...")
        
        # We run these effectively "in parallel" by batching requests (simulated here via sequential for stability, but high volume)
        board_members = {
            "Growth Officer": "Identify the #1 action to get revenue THIS WEEK.",
            "Risk Officer": "Analyze current plan for failures. What is the biggest risk?",
            "Chief Architect": "What is the most robust technical next step?"
        }
        
        board_outputs = {}
        for role, prompt in board_members.items():
            board_outputs[role] = self.spawn_agent(role, f"Context: Active Task is competitive intelligence. {prompt}")

        # Synthesize
        synthesis = self.spawn_agent("CEO", f"Review these Board reports: {board_outputs}. Decide the ONE unified goal for the next hour.")
        print(f"  👑 [CEO] Decision: {synthesis[:100]}...")

        # STAGE 2: THE FACTORY (5 Parallel Workers)
        print("  🏭 [Factory] Spinnning up 5 Worker Units...")
        
        # Manager breaks it down
        task_list_raw = self.spawn_agent("Factory Manager", f"Break this CEO goal into 5 distinct, isolated sub-tasks for workers: {synthesis}. Return ONLY a python list of strings.")
        
        try:
            # Safe eval to get list
            if "[" in task_list_raw:
                import ast
                # Extract list part
                start = task_list_raw.find("[")
                end = task_list_raw.rfind("]") + 1
                tasks = ast.literal_eval(task_list_raw[start:end])
            else:
                tasks = ["Review logs", "Check database", "Optimize prompts", "Verify backups", "Clean files"] # Fallback
                
            # Execute Workers
            worker_results = []
            for i, task in enumerate(tasks[:5]): # Cap at 5
                res = self.spawn_agent(f"Worker-{i+1}", f"Execute this task comprehensively: {task}")
                worker_results.append(res)
                
            print(f"  ✅ [Factory] production complete. {len(worker_results)} outputs generated.")
            
        except Exception as e:
            print(f"  ⚠️ [Factory] Assembly line error: {e}")

    def think(self, mode: str = "default"):
        """
        Autonomous thinking mode for unsupervised operation.
        Called during Unceasing Motion (night mode).
        """
        print(f"\n🌙 [Think] Entering autonomous mode: {mode}")

        if mode == "unsupervised_evolution":
            # Night mode: focus on background maintenance tasks
            tasks = [
                "Review and consolidate KNOWLEDGE_GRAPH entities for consistency",
                "Check for stale tasks in loop/tasks.json older than 7 days",
                "Validate all JSONL files in KNOWLEDGE_GRAPH/ have valid JSON",
                "Generate a summary of system health metrics",
                "Identify any orphaned files or unused imports in core/"
            ]

            import random
            selected_task = random.choice(tasks)
            print(f"  🎯 Selected task: {selected_task}")

            result = self.spawn_agent("Night Worker", selected_task, use_memory=True)
            print(f"  📝 Result: {result[:200]}...")

        else:
            # Default thinking mode
            self.run_evolution_cycle()

    def run_deep_design(self, task: str) -> str:
        """
        Trigger a Camel Inception Loop for complex design tasks.
        """
        from swarms.camel_tool import camel_tool
        return camel_tool.run_inception_loop(task)

if __name__ == "__main__":
    swarm = GenesisSwarm()
    # Test Camel Integration
    # print(swarm.run_deep_design("Design a distributed vector memory system"))
    
    # Default to High Velocity for acceleration
    swarm.run_high_velocity_cycle()