import time, uuid, hashlib, json, os, math
from enum import Enum
from .pheromones import PheromoneTrail

class OrganType(Enum):
    ANTENNAE = "antennae"
    SPINE = "spine"
    HIPPOCAMPUS = "hippocampus"
    CEREBELLUM = "cerebellum"
    PREFRONTAL = "prefrontal"
    DREAM = "dream"

ORGAN_DESCRIPTIONS = {
    OrganType.ANTENNAE: "Senses environment, classifies tasks, leaves pheromone trails",
    OrganType.SPINE: "Mamba-SSM sequence processor — pattern recognition & sequence modeling",
    OrganType.HIPPOCAMPUS: "Memory organ — stores and retrieves task embeddings",
    OrganType.CEREBELLUM: "Coordination organ — refines timing and motor sequences",
    OrganType.PREFRONTAL: "Planning organ — reasoning, decisions, multi-step plans",
    OrganType.DREAM: "Imagination organ — counterfactuals, simulation, what-if scenarios",
}

class OrganBase:
    def __init__(self, organ_type: OrganType, pheromones: PheromoneTrail):
        self.type = organ_type
        self.pheromones = pheromones
        self.id = str(uuid.uuid4())[:8]
        self.tasks_processed = 0
        self.tasks_succeeded = 0
        self.tasks_failed = 0
        self.log = []
        self.energy = 1.0

    def match_score(self, task: dict) -> float:
        return 0.0

    def process(self, task: dict) -> dict:
        raise NotImplementedError

    def _log(self, event: str, detail: dict = None):
        self.log.append({"time": time.time(), "organ": self.type.value, "event": event, "detail": detail or {}})

    def stats(self) -> dict:
        total = self.tasks_succeeded + self.tasks_failed
        return {
            "organ": self.type.value,
            "id": self.id,
            "description": ORGAN_DESCRIPTIONS.get(self.type, ""),
            "tasks_processed": self.tasks_processed,
            "tasks_succeeded": self.tasks_succeeded,
            "tasks_failed": self.tasks_failed,
            "success_rate": round(self.tasks_succeeded / total, 3) if total else 0,
            "energy": round(self.energy, 3),
            "log_entries": len(self.log),
        }

# ── Organ 1: Antennae ──
class AntennaeOrgan(OrganBase):
    def __init__(self, pheromones):
        super().__init__(OrganType.ANTENNAE, pheromones)
        self.task_keywords = {
            "code": ["build", "generate", "compile", "write", "create", "implement"],
            "memory": ["remember", "store", "recall", "dream", "save"],
            "plan": ["plan", "strategy", "design", "architecture", "roadmap"],
            "analyze": ["analyze", "audit", "scan", "check", "verify"],
            "imagine": ["imagine", "simulate", "what if", "suppose", "fantasize"],
            "search": ["search", "find", "hunt", "rabbit", "investigate"],
        }

    def _word_match(self, kw: str, text: str) -> bool:
        if " " in kw:
            return kw in text
        import re
        return bool(re.search(r'\b' + re.escape(kw) + r'\b', text))

    def classify(self, text: str) -> str:
        text_lower = text.lower()
        scores = {}
        for category, keywords in self.task_keywords.items():
            scores[category] = sum(1 for kw in keywords if self._word_match(kw, text_lower))
        if not scores or max(scores.values()) == 0:
            return "general"
        # Tie-breaking priority: imagine > memory > code > plan > analyze > search
        priority = ["imagine", "memory", "code", "plan", "analyze", "search"]
        max_score = max(scores.values())
        tied = [c for c, s in scores.items() if s == max_score]
        if len(tied) > 1:
            for p in priority:
                if p in tied:
                    return p
        return max(scores, key=scores.get)

    def match_score(self, task: dict) -> float:
        return 0.01

    def process(self, task: dict) -> dict:
        self.tasks_processed += 1
        text = task.get("input", "")
        task_type = self.classify(text)
        task_id = task.get("task_id", str(uuid.uuid4())[:8])
        self.pheromones.lay(task_id, "antennae", strength=0.8, metadata={"type": task_type})
        self._log("classified", {"task_id": task_id, "type": task_type})
        return {"task_id": task_id, "type": task_type, "organ": "antennae", "status": "classified"}

# ── Organ 2: Spine (Mamba-SSM) ──
class SpineOrgan(OrganBase):
    def __init__(self, pheromones):
        super().__init__(OrganType.SPINE, pheromones)
        self.mamba_available = False
        self._init_mamba()

    def _init_mamba(self):
        try:
            import mamba_ssm
            self.mamba_available = True
        except ImportError:
            self.mamba_available = False

    def _cpu_mamba_stub(self, sequence: list) -> dict:
        seq_str = " ".join(str(s) for s in sequence)
        seq_hash = hashlib.md5(seq_str.encode()).hexdigest()
        pattern_strength = 0.0
        for s in sequence:
            if isinstance(s, str):
                pattern_strength += len(s) / 100.0
            elif isinstance(s, (int, float)):
                pattern_strength += s / 1000.0
        return {
            "sequence_hash": seq_hash,
            "pattern_strength": min(1.0, pattern_strength),
            "sequence_length": len(sequence),
            "mamba_mode": "cpu_stub",
            "embeddings_dim": 128,
        }

    def match_score(self, task: dict) -> float:
        task_type = task.get("type", "")
        text = task.get("input", "").lower()
        if task_type in ("code", "analyze"):
            return 0.85
        if any(w in text for w in ["search", "find", "hunt", "evidence", "source"]):
            return 0.7
        return 0.3

    def process(self, task: dict) -> dict:
        self.tasks_processed += 1
        text = task.get("input", "")
        tokens = text.split()[:256]
        result = self._cpu_mamba_stub(tokens)
        task_id = task.get("task_id", "")
        self.pheromones.lay(task_id, "spine", strength=0.7, metadata={"mode": result["mamba_mode"]})
        self._log("processed", {"task_id": task_id, "tokens": len(tokens), "pattern": result["pattern_strength"]})
        return {"task_id": task_id, "organ": "spine", "status": "processed", "result": result}

# ── Organ 3: Hippocampus ──
class HippocampusOrgan(OrganBase):
    def __init__(self, pheromones):
        super().__init__(OrganType.HIPPOCAMPUS, pheromones)
        self.memories = []
        self.memory_file = "/tmp/fsi_felon/chimera/hippocampus_memory.json"
        self._load()

    def _load(self):
        if os.path.exists(self.memory_file):
            try:
                with open(self.memory_file) as f:
                    self.memories = json.load(f)
            except: pass

    def _save(self):
        try:
            with open(self.memory_file, 'w') as f:
                json.dump(self.memories[-500:], f, indent=2)
        except: pass

    def store(self, key: str, value: any, metadata: dict = None):
        entry = {"key": key, "value": value, "metadata": metadata or {}, "time": time.time()}
        self.memories.append(entry)
        self._save()
        return entry

    def recall(self, key: str) -> list:
        matches = [m for m in self.memories if key.lower() in m["key"].lower()]
        return matches[-5:]

    def match_score(self, task: dict) -> float:
        task_type = task.get("type", "")
        text = task.get("input", "").lower()
        if task_type == "memory":
            return 1.0
        if any(w in text for w in ["remember", "recall", "store", "save"]):
            return 0.9
        # Only match "dream" keyword if not also an imagine-type task
        if "dream" in text and not any(w in text for w in ["imagine", "what if", "suppose"]):
            return 0.6
        return 0.2

    def process(self, task: dict) -> dict:
        self.tasks_processed += 1
        text = task.get("input", "")
        task_id = task.get("task_id", "")
        if "store" in text.lower() or "save" in text.lower():
            parts = text.split(" ", 2)
            key = parts[-1] if len(parts) > 1 else "unnamed"
            self.store(key, text, {"source": "chimera"})
            result = {"action": "stored", "key": key}
        else:
            results = self.recall(text)
            result = {"action": "recalled", "matches": len(results), "memories": results[-3:] if results else []}
        self.pheromones.lay(task_id, "hippocampus", strength=0.6)
        self._log("processed", {"task_id": task_id, "action": result["action"]})
        return {"task_id": task_id, "organ": "hippocampus", "status": "processed", "result": result}

# ── Organ 4: Cerebellum ──
class CerebellumOrgan(OrganBase):
    def __init__(self, pheromones):
        super().__init__(OrganType.CEREBELLUM, pheromones)
        self.sequences = {}

    def learn_sequence(self, task_id: str, steps: list):
        self.sequences[task_id] = {"steps": steps, "learned": time.time(), "success_count": 0}

    def refine(self, task_id: str, feedback: dict) -> dict:
        if task_id in self.sequences:
            self.sequences[task_id]["success_count"] += 1
            return {"refined": True, "sequence_id": task_id, "confidence": min(1.0, self.sequences[task_id]["success_count"] * 0.2)}
        return {"refined": False, "reason": "unknown_sequence"}

    def match_score(self, task: dict) -> float:
        task_type = task.get("type", "")
        text = task.get("input", "").lower()
        import re
        if re.search(r'\bbuild\b', text):
            return 0.87
        if task_type == "code":
            return 0.8
        if any(w in text for w in ["search", "find", "hunt"]):
            return 0.3
        if task_type == "general":
            return 0.8
        return 0.5

    def process(self, task: dict) -> dict:
        self.tasks_processed += 1
        task_id = task.get("task_id", "")
        result = self.refine(task_id, {})
        self.pheromones.lay(task_id, "cerebellum", strength=0.5)
        self._log("processed", {"task_id": task_id, "refined": result["refined"]})
        return {"task_id": task_id, "organ": "cerebellum", "status": "processed", "result": result}

# ── Organ 5: Prefrontal ──
class PrefrontalOrgan(OrganBase):
    def __init__(self, pheromones):
        super().__init__(OrganType.PREFRONTAL, pheromones)
        self.plans = {}

    def plan(self, goal: str, constraints: list = None) -> dict:
        plan_id = str(uuid.uuid4())[:8]
        steps = [
            {"step": 1, "action": "analyze", "description": f"Analyze requirements for: {goal[:50]}"},
            {"step": 2, "action": "design", "description": "Design solution architecture"},
            {"step": 3, "action": "implement", "description": "Implement core components"},
            {"step": 4, "action": "verify", "description": "Verify correctness"},
            {"step": 5, "action": "deliver", "description": "Deliver result"},
        ]
        self.plans[plan_id] = {"goal": goal, "steps": steps, "created": time.time(), "constraints": constraints or []}
        return {"plan_id": plan_id, "steps": steps, "total_steps": len(steps)}

    def match_score(self, task: dict) -> float:
        task_type = task.get("type", "")
        if task_type == "plan":
            return 1.0
        if any(w in task.get("input", "").lower() for w in ["design", "architecture", "strategy", "plan"]):
            return 0.8
        return 0.3

    def process(self, task: dict) -> dict:
        self.tasks_processed += 1
        text = task.get("input", "")
        task_id = task.get("task_id", "")
        result = self.plan(text)
        self.pheromones.lay(task_id, "prefrontal", strength=0.9, metadata={"plan_id": result["plan_id"]})
        self._log("planned", {"task_id": task_id, "plan_id": result["plan_id"], "steps": result["total_steps"]})
        return {"task_id": task_id, "organ": "prefrontal", "status": "planned", "result": result}

# ── Organ 6: Dream ──
class DreamOrgan(OrganBase):
    def __init__(self, pheromones):
        super().__init__(OrganType.DREAM, pheromones)
        self.dreams = []

    def imagine(self, prompt: str, temperature: float = 0.85) -> dict:
        dream_id = str(uuid.uuid4())[:8]
        dream = {
            "id": dream_id,
            "prompt": prompt,
            "temperature": temperature,
            "content": self._generate_dream_content(prompt),
            "timestamp": time.time(),
            "energy_cost": round(temperature * 0.3 + 0.2, 3),
        }
        self.dreams.append(dream)
        return dream

    def _generate_dream_content(self, prompt: str) -> str:
        p = prompt.lower()
        if "code" in p or "build" in p:
            return f"Dream: I saw a structure of linked components handling {prompt[:30]}... Each module spoke to the next."
        if "truth" in p or "rabbit" in p:
            return f"Dream: Trails of evidence branching infinitely. Each path a different narrative. None complete alone."
        if "memory" in p or "remember" in p:
            return f"Dream: Fragments of past tasks floating in hyperbolic space. Connections forming between distant memories."
        return f"Dream: A landscape of possibility around '{prompt[:40]}'... Patterns emerging and collapsing."

    def match_score(self, task: dict) -> float:
        task_type = task.get("type", "")
        text = task.get("input", "").lower()
        if task_type == "imagine":
            return 1.0
        if any(w in text for w in ["imagine", "what if", "suppose", "fantasize"]):
            return 0.95
        if "dream" in text:
            return 0.7
        if self.energy > 0.5:
            return 0.2
        return 0.05

    def process(self, task: dict) -> dict:
        self.tasks_processed += 1
        text = task.get("input", "")
        task_id = task.get("task_id", "")
        dream = self.imagine(text)
        self.energy = max(0, self.energy - dream["energy_cost"])
        self.pheromones.lay(task_id, "dream", strength=dream["energy_cost"])
        self._log("dreamt", {"task_id": task_id, "dream_id": dream["id"], "energy_cost": dream["energy_cost"]})
        return {"task_id": task_id, "organ": "dream", "status": "processed", "result": dream}

ORGAN_MAP = {
    OrganType.ANTENNAE: AntennaeOrgan,
    OrganType.SPINE: SpineOrgan,
    OrganType.HIPPOCAMPUS: HippocampusOrgan,
    OrganType.CEREBELLUM: CerebellumOrgan,
    OrganType.PREFRONTAL: PrefrontalOrgan,
    OrganType.DREAM: DreamOrgan,
}

def create_all_organs(pheromones: PheromoneTrail) -> list:
    return [cls(pheromones) for cls in ORGAN_MAP.values()]