# queen_orchestrator.py
import time
import datetime
import json
import subprocess
import os
import logging
import redis
import psycopg2
import requests

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# Constants and Configuration
SPRINT_PLAN_PATH = "QUEEN_ELEVATION_SPRINT_PLAN.md"
GENESIS_PACKAGE_PATH = "GENESIS_COMPLETE_PACKAGE.md"
CHECKPOINTS_DIR = "sprint-checkpoints"
MODEL = "Qwen-Long 1.5 (30B parameters)"
OLLAMA_ENDPOINT = "http://152.53.201.152:23405"
REDIS_HOST = "redis-genesis-u50607.vm.elestio.app"
REDIS_PORT = 26379
REDIS_PASSWORD = ""  # Replace with actual password if needed
POSTGRES_HOST = "postgresql-genesis-u50607.vm.elestio.app"
POSTGRES_PORT = 5432
POSTGRES_USER = "postgres"
POSTGRES_PASSWORD = "" # Replace with actual password if needed
POSTGRES_DB = "genesis_memory"
QDRANT_HOST = "qdrant-b3knu-u50607.vm.elestio.app"
QDRANT_PORT = 6333
BUDGET_LIMIT = 10.00
EMERGENCY_STOP = 9.50
SPRINT_ID = "QUEEN-ELEVATION-2026-01-11"
TOTAL_AGENTS = 50

# Load Sprint Plan and Genesis Package
def load_sprint_plan(filepath):
    try:
        with open(filepath, "r") as f:
            content = f.read()
            # Basic parsing, improve as needed (e.g., regex, proper parser)
            plan = {}
            plan['phases'] = {}
            phase_lines = []
            phase_num = 0
            in_phase = False
            for line in content.splitlines():
                if line.startswith("## PHASE"):
                    if in_phase:
                        phase_data = parse_phase_data(phase_lines)
                        plan['phases'][f"Phase {phase_num}"] = phase_data
                        phase_lines = []
                    in_phase = True
                    phase_num = int(line.split(':')[0].split(' ')[2])
                    continue
                if in_phase:
                    phase_lines.append(line)
            if in_phase:
                phase_data = parse_phase_data(phase_lines)
                plan['phases'][f"Phase {phase_num}"] = phase_data
            return plan
    except FileNotFoundError:
        logging.error(f"Sprint plan file not found: {filepath}")
        return None

def parse_phase_data(lines):
    phase_data = {}
    agent_table_start = False
    agent_data = []
    for line in lines:
        if "Agent ID" in line and "Role" in line:
            agent_table_start = True
            continue
        if agent_table_start and "|" in line:
            parts = [p.strip() for p in line.split("|") if p.strip()]
            if len(parts) > 1:
                agent_data.append(parts)
    phase_data['agents'] = agent_data
    return phase_data

def load_genesis_package(filepath):
    try:
        with open(filepath, "r") as f:
            return f.read()
    except FileNotFoundError:
        logging.error(f"Genesis package file not found: {filepath}")
        return None

# Subsystem Management
def start_subsystem(subsystem_name):
    try:
        logging.info(f"Starting subsystem: {subsystem_name}")
        subprocess.run(["systemctl", "start", subsystem_name], check=True)
        logging.info(f"Subsystem {subsystem_name} started successfully.")
        return True
    except subprocess.CalledProcessError as e:
        logging.error(f"Failed to start subsystem {subsystem_name}: {e}")
        return False

def stop_subsystem(subsystem_name):
    try:
        logging.info(f"Stopping subsystem: {subsystem_name}")
        subprocess.run(["systemctl", "stop", subsystem_name], check=True)
        logging.info(f"Subsystem {subsystem_name} stopped successfully.")
        return True
    except subprocess.CalledProcessError as e:
        logging.error(f"Failed to stop subsystem {subsystem_name}: {e}")
        return False

# Health Monitoring
def check_redis_health():
    try:
        r = redis.Redis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, decode_responses=True)
        r.ping()
        logging.info("Redis health check passed.")
        return True
    except redis.exceptions.ConnectionError as e:
        logging.error(f"Redis health check failed: {e}")
        return False

def check_ollama_health():
    try:
        response = requests.get(f"{OLLAMA_ENDPOINT}/api/tags")
        response.raise_for_status()  # Raise HTTPError for bad responses (4xx or 5xx)
        logging.info("Ollama health check passed.")
        return True
    except requests.exceptions.RequestException as e:
        logging.error(f"Ollama health check failed: {e}")
        return False

def check_postgres_health():
    try:
        conn = psycopg2.connect(host=POSTGRES_HOST, port=POSTGRES_PORT, user=POSTGRES_USER, password=POSTGRES_PASSWORD, database=POSTGRES_DB)
        cur = conn.cursor()
        cur.execute("SELECT 1")
        result = cur.fetchone()
        conn.close()
        if result and result[0] == 1:
            logging.info("PostgreSQL health check passed.")
            return True
        else:
            logging.error("PostgreSQL health check failed: Unable to execute a simple query.")
            return False
    except psycopg2.Error as e:
        logging.error(f"PostgreSQL health check failed: {e}")
        return False

def check_qdrant_health():
    try:
        response = requests.get(f"http://{QDRANT_HOST}:{QDRANT_PORT}/collections/genesis_knowledge")
        response.raise_for_status()
        logging.info("Qdrant health check passed.")
        return True
    except requests.exceptions.RequestException as e:
        logging.error(f"Qdrant health check failed: {e}")
        return False

def run_health_checks():
    return {
        "redis": check_redis_health(),
        "ollama": check_ollama_health(),
        "postgres": check_postgres_health(),
        "qdrant": check_qdrant_health()
    }

# Resource Allocation (Token Budget)
class TokenBudgetMonitor:
    def __init__(self):
        self.BUDGET_LIMIT = BUDGET_LIMIT
        self.EMERGENCY_STOP = EMERGENCY_STOP
        self.current_spend = 0.0
        self.start_time = time.time()

    def update_spend(self, cost):
        self.current_spend += cost
        logging.info(f"Current spend: ${self.current_spend:.2f}")

    def check_budget(self):
        if self.current_spend >= self.EMERGENCY_STOP:
            logging.warning("Emergency stop threshold reached! Initiating graceful shutdown.")
            self.trigger_graceful_shutdown()
            return False
        elif self.current_spend >= self.BUDGET_LIMIT:
            logging.warning("Budget limit reached! Initiating graceful shutdown.")
            self.trigger_graceful_shutdown()
            return False
        return True

    def trigger_graceful_shutdown(self):
        logging.critical("Initiating graceful shutdown due to budget constraints.")
        # Implement graceful shutdown logic here (e.g., stop agents, save state)
        # This is a placeholder
        self.save_checkpoint("budget_shutdown")
        self.stop_all_agents()
        exit()

    def calculate_hourly_rate(self):
        elapsed_time = time.time() - self.start_time
        if elapsed_time > 0:
            return self.current_spend / (elapsed_time / 3600)
        else:
            return 0

    def generate_status_log(self, phase, hour, active_agents, completed_agents, pending_agents, failed_agents, checkpoints):
        now = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        spent_percent = (self.current_spend / self.BUDGET_LIMIT) * 100
        remaining = self.BUDGET_LIMIT - self.current_spend
        hourly_rate = self.calculate_hourly_rate()
        # Dummy token counts - replace with actual values
        input_tokens = 20000000
        output_tokens = 5000000
        total_tokens = input_tokens + output_tokens
        tokens_percent = (total_tokens / 53750000) * 100

        log = f"""
╔══════════════════════════════════════════════════════════════════╗
║                    AIVA QUEEN SPRINT - TOKEN STATUS              ║
╠══════════════════════════════════════════════════════════════════╣
║ Time: {now} | Phase: {phase}/5 | Hour: {hour:.1f}/10            ║
╠══════════════════════════════════════════════════════════════════╣
║ BUDGET                                                           ║
║ ├── Total: ${self.BUDGET_LIMIT:.2f}                                                ║
║ ├── Spent: ${self.current_spend:.2f} ({spent_percent:.1f}%)                                         ║
║ ├── Remaining: ${remaining:.2f}                                             ║
║ └── Rate: ${hourly_rate:.2f}/hr                                               ║
╠══════════════════════════════════════════════════════════════════╣
║ TOKENS                                                           ║
║ ├── Input:  {input_tokens/1000000:.1f}M tokens consumed                                ║
║ ├── Output: {output_tokens/1000000:.1f}M tokens generated                                ║
║ └── Total:  {total_tokens/1000000:.1f}M / 53.75M ({tokens_percent:.1f}%)                              ║
╠══════════════════════════════════════════════════════════════════╣
║ AGENTS                                                           ║
║ ├── Active: {active_agents}/{TOTAL_AGENTS}                                                ║
║ ├── Completed: {completed_agents}/{TOTAL_AGENTS}                                             ║
║ ├── Pending: {pending_agents}/{TOTAL_AGENTS}                                             ║
║ └── Failed: {failed_agents}/{TOTAL_AGENTS}                                                 ║
╠══════════════════════════════════════════════════════════════════╣
║ CHECKPOINTS                                                      ║"""
        for checkpoint_name, checkpoint_status in checkpoints.items():
            log += f"\n║ ├── [{checkpoint_status}] {checkpoint_name}"
        log += """
╚══════════════════════════════════════════════════════════════════╝
"""
        logging.info(log)

# Checkpoint Management
def save_checkpoint(checkpoint_name, data):
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"{CHECKPOINTS_DIR}/{SPRINT_ID}_{checkpoint_name}_{timestamp}.json"
    os.makedirs(CHECKPOINTS_DIR, exist_ok=True)
    try:
        with open(filename, "w") as f:
            json.dump(data, f, indent=4)
        logging.info(f"Checkpoint saved to: {filename}")
        return filename
    except Exception as e:
        logging.error(f"Failed to save checkpoint: {e}")
        return None

def load_checkpoint(checkpoint_file):
    try:
        with open(checkpoint_file, "r") as f:
            data = json.load(f)
        logging.info(f"Checkpoint loaded from: {checkpoint_file}")
        return data
    except FileNotFoundError:
        logging.error(f"Checkpoint file not found: {checkpoint_file}")
        return None
    except json.JSONDecodeError:
        logging.error(f"Invalid JSON in checkpoint file: {checkpoint_file}")
        return None

# Task Routing (Placeholder)
def route_task(agent_id, task_description):
    logging.info(f"Routing task for agent {agent_id}: {task_description}")
    # Implement task routing logic here, e.g., dispatch to a queue, call an API
    # This is a placeholder
    return f"Task '{task_description}' assigned to agent {agent_id}."

# Agent Management (Placeholder)
def start_agent(agent_id):
    logging.info(f"Starting agent: {agent_id}")
    # Implement agent startup logic here, e.g., spawn a process, call an API
    # This is a placeholder
    return f"Agent {agent_id} started."

def stop_agent(agent_id):
    logging.info(f"Stopping agent: {agent_id}")
    # Implement agent shutdown logic here, e.g., terminate a process, call an API
    # This is a placeholder
    return f"Agent {agent_id} stopped."

def stop_all_agents():
    logging.info("Stopping all agents...")
    # Placeholder - replace with actual agent shutdown logic
    logging.warning("Stopping all agents - IMPLEMENTATION NEEDED!")

# Strategic Decision-Making (Placeholder)
def make_strategic_decision(context):
    logging.info(f"Making strategic decision with context: {context}")
    # Implement strategic decision-making logic here, e.g., use Ollama to generate a plan
    # This is a placeholder
    return "Strategic decision: Proceed with Phase 2."

# Main Orchestration Function
class QueenOrchestrator:
    def __init__(self):
        self.sprint_plan = load_sprint_plan(SPRINT_PLAN_PATH)
        self.genesis_package = load_genesis_package(GENESIS_PACKAGE_PATH)
        self.budget_monitor = TokenBudgetMonitor()
        self.current_phase = 0
        self.start_time = time.time()
        self.active_agents = 0
        self.completed_agents = 0
        self.pending_agents = TOTAL_AGENTS
        self.failed_agents = 0
        self.checkpoints = {
            "Phase 1: Foundation (Hour 2)": " ",
            "Phase 2: Knowledge (Hour 4)": " ",
            "Phase 3: Capabilities (Hour 6)": " ",
            "Phase 4: Swarm (Hour 8)": " ",
            "Phase 5: Coronation (Hour 10)": " "
        }

    def execute_phase(self, phase_number):
        self.current_phase = phase_number
        phase_name = f"Phase {phase_number}"
        logging.info(f"Starting {phase_name}")
        if phase_name not in self.sprint_plan['phases']:
            logging.error(f"Phase {phase_name} not found in sprint plan.")
            return False

        phase_data = self.sprint_plan['phases'][phase_name]
        agent_tasks = phase_data.get('agents')
        if not agent_tasks:
            logging.warning(f"No agents defined for {phase_name}.")
            return True # Consider it successful if no agents to run

        num_agents = len(agent_tasks)
        self.active_agents += num_agents
        self.pending_agents -= num_agents
        logging.info(f"Starting {num_agents} agents for {phase_name}")

        for agent_data in agent_tasks:
            agent_id = agent_data[0]
            task_description = agent_data[2]  # Assuming task description is always the third element
            logging.info(f"Starting agent {agent_id} with task: {task_description}")
            route_task(agent_id, task_description) # Assign task to agent
            start_agent(agent_id)  # Start the agent

        # Simulate task completion (replace with actual agent completion signals)
        time.sleep(1)
        self.completed_agents += num_agents
        self.active_agents -= num_agents
        logging.info(f"All agents completed tasks for {phase_name}")
        self.checkpoints[phase_name + " (Hour "+str(phase_number*2)+")"] = "✓"
        return True

    def run(self):
        logging.info("Queen Orchestrator started.")
        health_checks = run_health_checks()
        logging.info(f"Initial health checks: {health_checks}")

        if not all(health_checks.values()):
            logging.critical("Critical system health checks failed. Aborting sprint.")
            return

        for phase_number in range(1, 6):
            hour = (time.time() - self.start_time) / 3600
            self.budget_monitor.generate_status_log(phase_number, hour, self.active_agents, self.completed_agents, self.pending_agents, self.failed_agents, self.checkpoints)
            if not self.budget_monitor.check_budget():
                logging.critical("Budget exceeded. Aborting sprint.")
                return
            if not self.execute_phase(phase_number):
                logging.error(f"Phase {phase_number} failed. Aborting sprint.")
                return
            hour = (time.time() - self.start_time) / 3600
            self.budget_monitor.generate_status_log(phase_number, hour, self.active_agents, self.completed_agents, self.pending_agents, self.failed_agents, self.checkpoints)
            self.save_checkpoint(f"phase_{phase_number}_complete", {"status": "success"}) # Save Checkpoint
            logging.info(f"Phase {phase_number} completed successfully.")
            time.sleep(1)  # Simulate processing time
            if phase_number == 1:
                self.checkpoints["Phase 1: Foundation (Hour 2)"] = "✓"
            elif phase_number == 2:
                self.checkpoints["Phase 2: Knowledge (Hour 4)"] = "✓"
            elif phase_number == 3:
                self.checkpoints["Phase 3: Capabilities (Hour 6)"] = "✓"
            elif phase_number == 4:
                self.checkpoints["Phase 4: Swarm (Hour 8)"] = "✓"
            elif phase_number == 5:
                self.checkpoints["Phase 5: Coronation (Hour 10)"] = "✓"

        logging.info("Queen Elevation Sprint completed successfully!")
        self.budget_monitor.generate_status_log(5, 10, self.active_agents, self.completed_agents, self.pending_agents, self.failed_agents, self.checkpoints)

if __name__ == "__main__":
    orchestrator = QueenOrchestrator()
    orchestrator.run()