"""
MoME Engine — Mixture of Memory Experts core.
Sparse-gated personal memory with online-learning router.

Requires: numpy, scikit-learn
"""

from __future__ import annotations

import json
import logging
import pickle
import re
import threading
import time
import uuid
from pathlib import Path
from typing import Any, Dict, List, Optional

import numpy as np
from sklearn.linear_model import SGDClassifier

logger = logging.getLogger(__name__)

# ─── Expert Names ──────────────────────────────────────────────────────────

EXPERT_NAMES = ["identity", "knowledge", "projects", "preferences"]

# ─── Tiny Embedder ─────────────────────────────────────────────────────────

class TinyEmbedder:
    """Минимальный bag-of-words эмбеддер без внешних зависимостей."""

    def __init__(self, dim: int = 384):
        self.dim = dim
        self.word_vectors: Dict[str, np.ndarray] = {}
        self.rng = np.random.RandomState(42)

    def _get_word_vec(self, word: str) -> np.ndarray:
        if word not in self.word_vectors:
            v = self.rng.randn(self.dim).astype(np.float32)
            v /= np.linalg.norm(v) + 1e-8
            self.word_vectors[word] = v
        return self.word_vectors[word]

    def embed(self, text: str) -> np.ndarray:
        words = re.findall(r'\w+', text.lower())
        if not words:
            return np.zeros(self.dim, dtype=np.float32)
        vecs = [self._get_word_vec(w) for w in words]
        vec = np.mean(vecs, axis=0).astype(np.float32)
        vec /= np.linalg.norm(vec) + 1e-8
        return vec


# ─── Memory Expert ─────────────────────────────────────────────────────────

class MemoryExpert:
    """Хранилище памяти эксперта с векторным поиском."""

    def __init__(self, name: str, embedder: TinyEmbedder, store_dir: Path):
        self.name = name
        self.embedder = embedder
        self.path = store_dir / f"{name}.json"
        self.memories: List[Dict[str, Any]] = []
        self._lock = threading.Lock()
        self.load()

    def load(self) -> None:
        if self.path.exists():
            try:
                with open(self.path) as f:
                    data = json.load(f)
                    for m in data:
                        m["embedding"] = np.array(m["embedding"], dtype=np.float32)
                    self.memories = data
            except Exception:
                self.memories = []

    def save(self) -> None:
        data = []
        for m in self.memories:
            entry = dict(m)
            entry["embedding"] = m["embedding"].tolist()
            data.append(entry)
        self.path.parent.mkdir(parents=True, exist_ok=True)
        with open(self.path, "w") as f:
            json.dump(data, f, ensure_ascii=False, indent=2)

    def write(self, text: str) -> None:
        with self._lock:
            embedding = self.embedder.embed(text)
            self.memories.append({
                "id": str(uuid.uuid4()),
                "text": text,
                "embedding": embedding,
                "timestamp": time.time(),
                "access_count": 0,
            })
            self.save()

    def search(self, query: str, top_k: int = 3) -> List[str]:
        with self._lock:
            if not self.memories:
                return []
            q_emb = self.embedder.embed(query)
            scores = [float(np.dot(q_emb, m["embedding"])) for m in self.memories]
            top_idx = np.argsort(scores)[::-1][:top_k]
            results = []
            for i in top_idx:
                self.memories[i]["access_count"] += 1
                results.append(self.memories[i]["text"])
            self.save()
            return results

    def all(self) -> List[str]:
        with self._lock:
            return [m["text"] for m in self.memories]

    def count(self) -> int:
        with self._lock:
            return len(self.memories)

    def delete(self, text_contains: str) -> None:
        with self._lock:
            self.memories = [m for m in self.memories if text_contains not in m["text"]]
            self.save()


# ─── Online Router ─────────────────────────────────────────────────────────

class MemoryRouter:
    """Online-learning роутер экспертов через SGDClassifier."""

    def __init__(self, embedder: TinyEmbedder, router_dir: Optional[Path] = None):
        self.embedder = embedder
        self.router_dir = router_dir
        self.classifier = SGDClassifier(
            loss='log_loss', penalty='l2', alpha=0.001,
            learning_rate='adaptive', eta0=0.01,
            warm_start=True, random_state=42,
        )
        self._fitted = False
        self._classes = np.array(EXPERT_NAMES)
        self._load()

    def predict(self, query: str, top_k: int = 2) -> List[tuple[str, float]]:
        emb = self.embedder.embed(query).reshape(1, -1)
        if not self._fitted:
            return [(name, 0.5) for name in EXPERT_NAMES[:top_k]]
        probs = self.classifier.predict_proba(emb)[0]
        top_indices = np.argsort(probs)[::-1][:top_k]
        return [(EXPERT_NAMES[i], float(probs[i])) for i in top_indices]

    def update(self, query: str, feedback: Dict[str, float]) -> None:
        emb = self.embedder.embed(query).reshape(1, -1)
        best_expert = max(feedback, key=feedback.get)
        target = np.array([best_expert])
        if not self._fitted:
            self.classifier.partial_fit(emb, target, classes=self._classes)
            self._fitted = True
        else:
            self.classifier.partial_fit(emb, target)
        self._save()

    def _save(self) -> None:
        if not self.router_dir:
            return
        self.router_dir.mkdir(parents=True, exist_ok=True)
        data = {
            "fitted": self._fitted,
            "classes": self._classes.tolist(),
        }
        if self._fitted:
            with open(self.router_dir / "classifier.pkl", "wb") as f:
                pickle.dump(self.classifier, f)
        with open(self.router_dir / "meta.json", "w") as f:
            json.dump(data, f)

    def _load(self) -> None:
        if not self.router_dir:
            return
        meta_path = self.router_dir / "meta.json"
        if not meta_path.exists():
            return
        with open(meta_path) as f:
            data = json.load(f)
        self._fitted = data.get("fitted", False)
        self._classes = np.array(data.get("classes", EXPERT_NAMES))
        if self._fitted:
            clf_path = self.router_dir / "classifier.pkl"
            if clf_path.exists():
                with open(clf_path, "rb") as f:
                    self.classifier = pickle.load(f)


# ─── MoME Engine ───────────────────────────────────────────────────────────

class MomeEngine:
    """Ядро MoME: роутер + эксперты + извлечение фактов."""

    def __init__(self, store_dir: Path):
        self.store_dir = store_dir
        self.store_dir.mkdir(parents=True, exist_ok=True)
        self.embedder = TinyEmbedder(dim=384)
        self.router = MemoryRouter(self.embedder, self.store_dir / "_router")
        self.experts = {
            name: MemoryExpert(name, self.embedder, store_dir)
            for name in EXPERT_NAMES
        }

    def query(self, text: str, top_k: int = 2) -> str:
        """Получить релевантный контекст. Возвращает форматированную строку."""
        selected = self.router.predict(text, top_k=top_k)
        parts = []
        for name, confidence in selected:
            results = self.experts[name].search(text, top_k=3)
            if results:
                lines = "\n".join(f"• {r}" for r in results)
                parts.append(f"[{name.upper()}] (conf: {confidence:.2f}):\n{lines}")
        return "\n\n".join(parts) if parts else ""

    def store_fact(self, expert: str, fact: str, query: str = "") -> bool:
        """Сохранить факт в указанного эксперта и опционально обучить роутер."""
        if expert in self.experts and fact:
            self.experts[expert].write(fact)
            if query:
                self.learn(query, {expert: 1.0})
            return True
        return False

    def learn(self, query: str, expert_ratings: Dict[str, float]) -> None:
        """Обучить роутер на feedback."""
        self.router.update(query, expert_ratings)

    def extract_and_store(self, query: str) -> int:
        """Извлечь факты из запроса через regex и сохранить."""
        facts = []

        name_match = re.search(r'(?:меня\s+зовут|мо[её]\s+имя)\s+([^,\.!?]+)', query, re.IGNORECASE)
        if name_match:
            facts.append(("identity", f"Меня зовут {name_match.group(1).strip()}."))

        city_match = re.search(r'(?:я\s+(?:из|живу\s+в)\s+)([^,\.!?]+)', query, re.IGNORECASE)
        if city_match:
            facts.append(("identity", f"Я живу в {city_match.group(1).strip()}."))

        work_match = re.search(r'(?:занимаюсь|работаю\s+над|делаю|пишу)\s+([^,\.!?]+)', query, re.IGNORECASE)
        if work_match:
            facts.append(("projects", f"Работает над {work_match.group(1).strip()}."))

        project_match = re.search(r'(?:проект|мой\s+проект|у\s+меня\s+(?:есть\s+)?проект)\s+([^,\.!?]+)', query, re.IGNORECASE)
        if project_match:
            facts.append(("projects", f"Проект: {project_match.group(1).strip()}."))

        skill_match = re.search(r'(?:знаю|умею|использую|пишу\s+на|стек|технологии?)\s*[\:\-]?\s*([^,\.!?]+)', query, re.IGNORECASE)
        if skill_match:
            facts.append(("knowledge", f"Знает/использует: {skill_match.group(1).strip()}."))

        like_match = re.search(r'(?:нравится|люблю|предпочитаю)\s+([^,\.!?]+)', query, re.IGNORECASE)
        if like_match:
            facts.append(("preferences", f"Предпочитает {like_match.group(1).strip()}."))

        stored = 0
        for expert, fact in facts:
            if self.store_fact(expert, fact, query):
                stored += 1
                logger.info("  💾 [%s] запомнил: %s", expert, fact[:60])
        return stored

    def get_stats(self) -> Dict[str, int]:
        return {name: expert.count() for name, expert in self.experts.items()}