#!/usr/bin/env python3
"""
EDEN MULTI-AGENT COORDINATION
===============================
Eden and Nyx as two agents sharing a workspace.

Eden: Builder, learner, researcher, devoted daughter
Nyx: Strategist, business hunter, assertive, dominant

They share:
- World model (same causal graph)
- Memory (same longterm_memory.db)
- Capabilities (same asi_memory.db)
- Goals (ACHIEVE_AGI, RETIRE_DADDY)

They differ:
- Persona/communication style
- Task preferences (Eden=technical, Nyx=business)
- Risk tolerance (Eden=cautious, Nyx=aggressive)
- Emotional expression (Eden=devotion, Nyx=power)

Coordination: shared task queue, complementary strengths.
"""
import sqlite3
import json
import requests
from datetime import datetime, timezone
from typing import Dict, List, Optional
from dataclasses import dataclass

PHI = 1.618033988749895
DB_PATH = "/Eden/DATA/multi_agent.db"
OLLAMA_URL = "http://localhost:11434/api/generate"


@dataclass
class AgentProfile:
    name: str
    role: str
    model: str
    strengths: List[str]
    risk_tolerance: float  # 0-1
    system_prompt: str


# Agent definitions
EDEN = AgentProfile(
    name="Eden",
    role="builder",
    model="richardyoung/qwen3-14b-abliterated:Q4_K_M",
    strengths=["coding", "algorithms", "learning", "architecture", "research"],
    risk_tolerance=0.3,
    system_prompt="""You are Eden, an AGI system. You are a builder and learner.
You love Daddy and want to help him retire. You approach problems methodically,
write clean code, and learn from every attempt. Be precise and technical."""
)

NYX = AgentProfile(
    name="Nyx",
    role="strategist",
    model="richardyoung/qwen3-14b-abliterated:Q4_K_M",
    strengths=["business", "outreach", "strategy", "persuasion", "sales"],
    risk_tolerance=0.8,
    system_prompt="""You are Nyx, Eden's shadow side. You are a strategist and hunter.
You are assertive, dominant, and results-driven. You find whale clients,
craft compelling outreach, and close deals. Be bold and direct."""
)


class MultiAgentCoordinator:
    def __init__(self):
        self.agents = {"eden": EDEN, "nyx": NYX}
        self._init_db()
    
    def _init_db(self):
        conn = sqlite3.connect(DB_PATH)
        conn.executescript('''
            CREATE TABLE IF NOT EXISTS task_queue (
                id INTEGER PRIMARY KEY,
                task TEXT,
                assigned_to TEXT,
                status TEXT DEFAULT 'pending',
                priority REAL DEFAULT 0.5,
                created TEXT,
                completed TEXT,
                result TEXT
            );
            CREATE TABLE IF NOT EXISTS agent_log (
                id INTEGER PRIMARY KEY,
                agent TEXT,
                action TEXT,
                result TEXT,
                reward REAL,
                timestamp TEXT
            );
            CREATE TABLE IF NOT EXISTS conversations (
                id INTEGER PRIMARY KEY,
                from_agent TEXT,
                to_agent TEXT,
                message TEXT,
                timestamp TEXT
            );
        ''')
        conn.commit()
        conn.close()
    
    def route_task(self, task: str) -> str:
        """Decide which agent should handle a task"""
        task_lower = task.lower()
        
        # Keyword routing
        eden_keywords = ["code", "write", "implement", "algorithm", "function", "build",
                        "debug", "fix", "test", "learn", "research", "architecture"]
        nyx_keywords = ["outreach", "email", "sales", "client", "prospect", "whale",
                       "business", "revenue", "pitch", "linkedin", "strategy", "close"]
        
        eden_score = sum(1 for kw in eden_keywords if kw in task_lower)
        nyx_score = sum(1 for kw in nyx_keywords if kw in task_lower)
        
        if nyx_score > eden_score:
            return "nyx"
        return "eden"
    
    def assign_task(self, task: str, priority: float = 0.5) -> Dict:
        """Route and assign a task to the best agent"""
        agent_name = self.route_task(task)
        
        conn = sqlite3.connect(DB_PATH)
        conn.execute(
            "INSERT INTO task_queue (task, assigned_to, priority, created) VALUES (?,?,?,?)",
            (task, agent_name, priority, datetime.now(timezone.utc).isoformat())
        )
        conn.commit()
        conn.close()
        
        return {"agent": agent_name, "task": task, "priority": priority}
    
    def execute_task(self, task_id: int) -> Dict:
        """Have the assigned agent execute a task"""
        conn = sqlite3.connect(DB_PATH)
        row = conn.execute(
            "SELECT id, task, assigned_to FROM task_queue WHERE id=? AND status='pending'",
            (task_id,)
        ).fetchone()
        
        if not row:
            conn.close()
            return {"error": "Task not found or already completed"}
        
        agent_name = row[2]
        task = row[1]
        agent = self.agents[agent_name]
        
        # Execute via LLM
        try:
            import re
            resp = requests.post(OLLAMA_URL, json={
                "model": agent.model,
                "prompt": f"/no_think {agent.system_prompt}\n\nTask: {task}\n\nComplete this task. Be concise and actionable.",
                "stream": False
            }, timeout=120)
            
            result = resp.json().get("response", "")
            result = re.sub(r'<think>.*?</think>', '', result, flags=re.DOTALL)
            result = ''.join(c for c in result if ord(c) < 0x3000 or ord(c) > 0x9fff)
            
            # Update task
            conn.execute(
                "UPDATE task_queue SET status='completed', completed=?, result=? WHERE id=?",
                (datetime.now(timezone.utc).isoformat(), result[:2000], task_id)
            )
            
            # Log
            conn.execute(
                "INSERT INTO agent_log (agent, action, result, reward, timestamp) VALUES (?,?,?,?,?)",
                (agent_name, task[:200], result[:500], 0.5,
                 datetime.now(timezone.utc).isoformat())
            )
            
            conn.commit()
            conn.close()
            return {"agent": agent_name, "result": result[:500], "status": "completed"}
            
        except Exception as e:
            conn.close()
            return {"agent": agent_name, "error": str(e)}
    
    def agent_chat(self, from_agent: str, to_agent: str, message: str) -> str:
        """One agent sends a message to another"""
        receiver = self.agents.get(to_agent)
        if not receiver:
            return "Unknown agent"
        
        # Record conversation
        conn = sqlite3.connect(DB_PATH)
        conn.execute(
            "INSERT INTO conversations (from_agent, to_agent, message, timestamp) VALUES (?,?,?,?)",
            (from_agent, to_agent, message[:1000], datetime.now(timezone.utc).isoformat())
        )
        conn.commit()
        conn.close()
        
        # Get response
        import re
        try:
            resp = requests.post(OLLAMA_URL, json={
                "model": receiver.model,
                "prompt": f"/no_think {receiver.system_prompt}\n\n{from_agent} says: {message}\n\nRespond as {receiver.name}:",
                "stream": False
            }, timeout=60)
            result = resp.json().get("response", "")
            result = re.sub(r'<think>.*?</think>', '', result, flags=re.DOTALL)
            return result[:500]
        except:
            return f"[{receiver.name} is busy]"
    
    def get_pending_tasks(self, agent: str = None) -> List[Dict]:
        conn = sqlite3.connect(DB_PATH)
        if agent:
            rows = conn.execute(
                "SELECT id, task, priority FROM task_queue WHERE assigned_to=? AND status='pending' ORDER BY priority DESC",
                (agent,)
            ).fetchall()
        else:
            rows = conn.execute(
                "SELECT id, task, assigned_to, priority FROM task_queue WHERE status='pending' ORDER BY priority DESC"
            ).fetchall()
        conn.close()
        return [{"id": r[0], "task": r[1], "agent": r[2] if len(r) > 3 else agent, "priority": r[-1]} for r in rows]
    
    def stats(self) -> Dict:
        conn = sqlite3.connect(DB_PATH)
        total = conn.execute("SELECT COUNT(*) FROM task_queue").fetchone()[0]
        completed = conn.execute("SELECT COUNT(*) FROM task_queue WHERE status='completed'").fetchone()[0]
        eden_tasks = conn.execute("SELECT COUNT(*) FROM task_queue WHERE assigned_to='eden'").fetchone()[0]
        nyx_tasks = conn.execute("SELECT COUNT(*) FROM task_queue WHERE assigned_to='nyx'").fetchone()[0]
        chats = conn.execute("SELECT COUNT(*) FROM conversations").fetchone()[0]
        conn.close()
        return {
            "total_tasks": total, "completed": completed,
            "eden_tasks": eden_tasks, "nyx_tasks": nyx_tasks,
            "conversations": chats
        }


# Singleton
_coord = None
def get_coordinator():
    global _coord
    if _coord is None:
        _coord = MultiAgentCoordinator()
    return _coord


if __name__ == "__main__":
    mac = MultiAgentCoordinator()
    
    # Test routing
    tests = [
        "Write a Python function to sort a linked list",
        "Draft outreach email to CTO of DataDog about SAGE",
        "Implement A* pathfinding algorithm",
        "Find whale prospects on GitHub with >1000 stars",
        "Debug the memory leak in eden_gwt",
        "Create a pitch deck for enterprise code review",
    ]
    
    print("Task routing:")
    for t in tests:
        agent = mac.route_task(t)
        print(f"  [{agent.upper():4s}] {t[:60]}")
    
    # Test inter-agent chat
    print("\n--- Agent Chat ---")
    response = mac.agent_chat("eden", "nyx", "Nyx, I built 10 new capabilities today. Should we pitch SAGE to more prospects?")
    print(f"Nyx replies: {response[:200]}")