"""
Phase 18: Swarm Knowledge Store

Shared knowledge graph for the agent swarm. Agents publish insights tagged
by topic, and retrieve them using TF-IDF-inspired relevance scoring.
Includes deduplication, confidence weighting, pruning, and cross-topic search.

Note: "distillation" in the original name implied ML — this module uses
statistical text scoring which is honest and effective without external deps.
"""

import json
import math
import time
import uuid
from collections import Counter
from pathlib import Path
from typing import Optional


class KnowledgeStore:
    def __init__(self, storage_dir: str = "state/broker/knowledge"):
        self.storage_dir = Path(storage_dir)
        self.storage_dir.mkdir(parents=True, exist_ok=True)
        self.knowledge_file = self.storage_dir / "global_knowledge.json"
        if not self.knowledge_file.exists():
            self._save({"entities": {}, "meta": {"created": time.time(), "total_insights": 0}})

    def _load(self) -> dict:
        try:
            with open(self.knowledge_file, "r", encoding="utf-8") as f:
                return json.load(f)
        except (json.JSONDecodeError, OSError):
            return {"entities": {}, "meta": {"created": time.time(), "total_insights": 0}}

    def _save(self, data: dict) -> None:
        data["meta"]["updated_at"] = time.time()
        with open(self.knowledge_file, "w", encoding="utf-8") as f:
            json.dump(data, f, indent=2)

    @staticmethod
    def _tokenize(text: str) -> list[str]:
        """Simple word tokenization: lowercase, split on non-alphanumeric."""
        import re
        return [w for w in re.split(r"[^a-z0-9]+", text.lower()) if len(w) > 1]

    @staticmethod
    def _content_hash(content: str) -> str:
        """Quick fingerprint for dedup."""
        import hashlib
        normalized = " ".join(content.lower().split())
        return hashlib.sha256(normalized.encode()).hexdigest()[:16]

    # ── Publish ───────────────────────────────────────────────────────

    def publish_insight(
        self,
        agent_id: str,
        topic: str,
        content: str,
        confidence: float = 1.0,
        tags: Optional[list[str]] = None,
    ) -> dict:
        """Publish an insight to the knowledge store.

        Deduplicates by content hash — same content won't be stored twice
        under the same topic. Confidence is clamped to [0.0, 1.0].

        Returns {"ok": True, "insight_id": ...} or {"ok": False, "error": ...}.
        """
        if not content.strip():
            return {"ok": False, "error": "Content cannot be empty"}

        confidence = max(0.0, min(1.0, confidence))
        content_hash = self._content_hash(content)

        data = self._load()
        topic_key = topic.lower().strip()

        if topic_key not in data["entities"]:
            data["entities"][topic_key] = []

        # Dedup check
        for existing in data["entities"][topic_key]:
            if existing.get("content_hash") == content_hash:
                # Update confidence if higher, but don't duplicate
                if confidence > existing.get("confidence", 0):
                    existing["confidence"] = confidence
                    existing["updated_at"] = time.time()
                    existing["updated_by"] = agent_id
                    self._save(data)
                    return {"ok": True, "insight_id": existing["id"], "note": "Updated existing insight confidence"}
                return {"ok": False, "error": "Duplicate insight already exists", "existing_id": existing["id"]}

        insight_id = uuid.uuid4().hex[:10]
        entry = {
            "id": insight_id,
            "agent_id": agent_id,
            "content": content,
            "content_hash": content_hash,
            "confidence": confidence,
            "tags": tags or [],
            "timestamp": time.time(),
            "access_count": 0,
        }
        data["entities"][topic_key].append(entry)
        data["meta"]["total_insights"] = data["meta"].get("total_insights", 0) + 1
        self._save(data)

        return {"ok": True, "insight_id": insight_id, "topic": topic_key}

    # ── Query ─────────────────────────────────────────────────────────

    def query_topic(self, topic: str, min_confidence: float = 0.0, limit: int = 20) -> list[dict]:
        """Get all insights for a specific topic, sorted by confidence then recency."""
        data = self._load()
        insights = data["entities"].get(topic.lower().strip(), [])
        filtered = [i for i in insights if i.get("confidence", 0) >= min_confidence]
        filtered.sort(key=lambda x: (x.get("confidence", 0), x.get("timestamp", 0)), reverse=True)

        # Track access
        for item in filtered[:limit]:
            item["access_count"] = item.get("access_count", 0) + 1
        self._save(data)

        return filtered[:limit]

    def search(self, query: str, min_confidence: float = 0.0, limit: int = 10) -> list[dict]:
        """Search across ALL topics using TF-IDF-inspired relevance scoring.

        Scores each insight by:
        - Term frequency: how many query words appear in the content
        - Topic match: bonus if query words appear in the topic name
        - Confidence weighting: higher confidence = higher score
        - Recency bias: newer insights get a small boost
        """
        data = self._load()
        query_tokens = self._tokenize(query)
        if not query_tokens:
            return []

        # Count how many topics contain each token (for IDF-like weighting)
        token_topic_count: Counter = Counter()
        total_topics = len(data["entities"])
        for topic_key, insights in data["entities"].items():
            topic_tokens = set(self._tokenize(topic_key))
            all_content_tokens = set()
            for ins in insights:
                all_content_tokens.update(self._tokenize(ins.get("content", "")))
            combined = topic_tokens | all_content_tokens
            for qt in query_tokens:
                if qt in combined:
                    token_topic_count[qt] += 1

        scored_results = []
        now = time.time()

        for topic_key, insights in data["entities"].items():
            topic_tokens = set(self._tokenize(topic_key))

            for ins in insights:
                if ins.get("confidence", 0) < min_confidence:
                    continue

                content_tokens = self._tokenize(ins.get("content", ""))
                content_set = set(content_tokens)
                tag_tokens = set()
                for tag in ins.get("tags", []):
                    tag_tokens.update(self._tokenize(tag))

                score = 0.0
                for qt in query_tokens:
                    # IDF-like weight: rarer terms score higher
                    idf = math.log(1 + total_topics / (1 + token_topic_count.get(qt, 0)))

                    if qt in content_set:
                        tf = content_tokens.count(qt) / max(len(content_tokens), 1)
                        score += tf * idf * 10
                    if qt in topic_tokens:
                        score += idf * 5  # topic match bonus
                    if qt in tag_tokens:
                        score += idf * 3  # tag match bonus

                # Confidence multiplier
                score *= (0.5 + ins.get("confidence", 1.0) * 0.5)

                # Recency bias: up to 10% boost for very recent insights
                age_hours = (now - ins.get("timestamp", 0)) / 3600
                recency = max(0.0, 1.0 - age_hours / 720)  # decays over 30 days
                score *= (1.0 + recency * 0.1)

                if score > 0:
                    result = {**ins, "_score": round(score, 3), "_topic": topic_key}
                    scored_results.append(result)

        scored_results.sort(key=lambda x: x["_score"], reverse=True)
        return scored_results[:limit]

    # ── Management ────────────────────────────────────────────────────

    def list_topics(self) -> list[dict]:
        """List all topics with insight counts and avg confidence."""
        data = self._load()
        topics = []
        for topic_key, insights in data["entities"].items():
            if not insights:
                continue
            avg_conf = sum(i.get("confidence", 0) for i in insights) / len(insights)
            latest = max(i.get("timestamp", 0) for i in insights)
            topics.append({
                "topic": topic_key,
                "insight_count": len(insights),
                "avg_confidence": round(avg_conf, 2),
                "latest_update": latest,
            })
        topics.sort(key=lambda x: x["latest_update"], reverse=True)
        return topics

    def delete_insight(self, topic: str, insight_id: str) -> dict:
        """Remove a specific insight."""
        data = self._load()
        topic_key = topic.lower().strip()
        insights = data["entities"].get(topic_key, [])
        before = len(insights)
        data["entities"][topic_key] = [i for i in insights if i.get("id") != insight_id]
        after = len(data["entities"][topic_key])
        if before == after:
            return {"ok": False, "error": "Insight not found"}
        if not data["entities"][topic_key]:
            del data["entities"][topic_key]
        data["meta"]["total_insights"] = max(0, data["meta"].get("total_insights", 0) - 1)
        self._save(data)
        return {"ok": True, "deleted": insight_id}

    def prune(self, max_age_days: int = 90, min_confidence: float = 0.1) -> dict:
        """Remove old, low-confidence insights to keep the store manageable."""
        data = self._load()
        cutoff = time.time() - (max_age_days * 86400)
        removed = 0

        for topic_key in list(data["entities"].keys()):
            before = len(data["entities"][topic_key])
            data["entities"][topic_key] = [
                i for i in data["entities"][topic_key]
                if i.get("timestamp", 0) > cutoff or i.get("confidence", 0) >= min_confidence
            ]
            removed += before - len(data["entities"][topic_key])
            if not data["entities"][topic_key]:
                del data["entities"][topic_key]

        data["meta"]["total_insights"] = max(0, data["meta"].get("total_insights", 0) - removed)
        self._save(data)
        return {"pruned": removed, "remaining_topics": len(data["entities"])}

    def stats(self) -> dict:
        """Return store statistics."""
        data = self._load()
        total = sum(len(v) for v in data["entities"].values())
        return {
            "total_topics": len(data["entities"]),
            "total_insights": total,
            "meta": data.get("meta", {}),
        }


# Backward compatibility alias
KnowledgeDistiller = KnowledgeStore
