import sqlite3
import json
import time
import numpy as np
from sentence_transformers import SentenceTransformer


class Memory:
    def __init__(
        self,
        db_path="memory.db",
        embedding_model="intfloat/multilingual-e5-large"
    ):
        self.db = sqlite3.connect(db_path)
        self.db.row_factory = sqlite3.Row

        self.model = SentenceTransformer(embedding_model)

        self._create_tables()
        self._setup_prototypes()

    # ---------------------------
    # DB SETUP
    # ---------------------------

    def _create_tables(self):
        self.db.execute("""
        CREATE TABLE IF NOT EXISTS memories(
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            role TEXT,
            content TEXT,
            timestamp REAL,
            importance REAL,
            recalls INTEGER DEFAULT 0,
            last_accessed REAL DEFAULT 0,
            embedding TEXT
        )
        """)
        self.db.commit()

    # ---------------------------
    # EMBEDDINGS
    # ---------------------------

    def embed(self, text):
        return self.model.encode(
            text,
            normalize_embeddings=True
        ).tolist()

    def cosine_similarity(self, a, b):
        a = np.array(a)
        b = np.array(b)
        return float(np.dot(a, b))  # normalized embeddings => dot = cosine

    # ---------------------------
    # PROTOTYPES (SEMANTIC TYPES)
    # ---------------------------

    def _setup_prototypes(self):

        self.prototype_categories = {
            "identity": [
                "My name is John",
                "I am 25 years old",
                "I live in Berlin",
                "I work as a teacher"
            ],
            "preferences": [
                "My favorite game is Minecraft",
                "I love pizza",
                "I prefer cats"
            ],
            "relationships": [
                "I have a sister",
                "My wife is a doctor",
                "My best friend is Alex"
            ],
            "goals": [
                "I want to learn Python",
                "I plan to move abroad",
                "I am saving money"
            ],
            "health": [
                "I am allergic to peanuts",
                "I have diabetes",
                "I take medication"
            ],
            "temporary": [
                "I ate pizza today",
                "I watched a movie",
                "The weather is nice"
            ]
        }

        self.prototype_embeddings = {}

        for cat, examples in self.prototype_categories.items():
            self.prototype_embeddings[cat] = self.model.encode(
                examples,
                normalize_embeddings=True
            )

    # ---------------------------
    # CATEGORY DETECTION
    # ---------------------------

    def detect_memory_category(self, text, embedding):
        best_category = "temporary"
        best_score = -1.0

        for category, prototypes in self.prototype_embeddings.items():

            sims = [
                self.cosine_similarity(embedding, p)
                for p in prototypes
            ]

            score = max(sims)

            if score > best_score:
                best_score = score
                best_category = category

        return best_category, best_score

    # ---------------------------
    # NOVELTY
    # ---------------------------

    def novelty_score(self, embedding):
        rows = self.db.execute(
            "SELECT embedding FROM memories"
        ).fetchall()

        if not rows:
            return 1.0

        max_sim = 0.0

        for row in rows:
            stored = json.loads(row["embedding"])
            sim = self.cosine_similarity(embedding, stored)
            if sim > max_sim:
                max_sim = sim

        return float(max(0.0, 1.0 - max_sim))

    # ---------------------------
    # DUPLICATE CHECK
    # ---------------------------

    def is_duplicate(self, embedding, threshold=0.92):
        rows = self.db.execute(
            "SELECT embedding FROM memories"
        ).fetchall()

        for row in rows:
            stored = json.loads(row["embedding"])
            sim = self.cosine_similarity(embedding, stored)
            if sim >= threshold:
                return True

        return False

    # ---------------------------
    # IMPORTANCE
    # ---------------------------

    def calculate_importance(self, text, embedding):

        category, confidence = self.detect_memory_category(
            text,
            embedding
        )

        novelty = self.novelty_score(embedding)

        weights = {
            "identity": 1.0,
            "health": 1.0,
            "relationships": 0.95,
            "goals": 0.9,
            "preferences": 0.75,
            "temporary": 0.2
        }

        semantic_importance = confidence * weights[category]

        importance = (
            semantic_importance * 0.7 +
            novelty * 0.3
        )

        return float(np.clip(importance, 0.0, 1.0))

    # ---------------------------
    # ADD MEMORY
    # ---------------------------

    def add(self, role, content):

        embedding = self.embed(content)

        if self.is_duplicate(embedding):
            return

        importance = self.calculate_importance(content, embedding)

        self.db.execute("""
        INSERT INTO memories(
            role,
            content,
            timestamp,
            importance,
            embedding
        )
        VALUES (?, ?, ?, ?, ?)
        """, (
            role,
            content,
            time.time(),
            importance,
            json.dumps(embedding)
        ))

        self.db.commit()

    # ---------------------------
    # RETRIEVAL
    # ---------------------------

    def retrieve(self, query, top_k=10):

        query_embedding = self.embed(query)

        rows = self.db.execute(
            "SELECT * FROM memories"
        ).fetchall()

        now = time.time()
        scored = []

        for row in rows:

            embedding = json.loads(row["embedding"])

            similarity = self.cosine_similarity(
                query_embedding,
                embedding
            )

            age_days = (now - row["timestamp"]) / 86400

            recency = 1 / (1 + age_days * 0.05)

            recall_bonus = min(row["recalls"] * 0.02, 0.2)

            final_score = (
                similarity
                * (1 + row["importance"])
                * (1 + recency * 0.3)
                * (1 + recall_bonus)
            )

            scored.append((final_score, row))

        scored.sort(key=lambda x: x[0], reverse=True)

        memories = []

        for _, row in scored[:top_k]:

            self.db.execute("""
            UPDATE memories
            SET recalls = recalls + 1,
                last_accessed = ?
            WHERE id = ?
            """, (now, row["id"]))

            memories.append({
                "id": row["id"],
                "role": row["role"],
                "content": row["content"],
                "importance": row["importance"]
            })

        self.db.commit()
        return memories

    # ---------------------------
    # CONTEXT BUILDER
    # ---------------------------

    def build_context(self, query, top_k=10):
        memories = self.retrieve(query, top_k)

        return "\n".join(
            f"{m['role']}: {m['content']}"
            for m in memories
        )

    # ---------------------------
    # UTILITIES
    # ---------------------------

    def recent(self, limit=20):
        rows = self.db.execute("""
        SELECT * FROM memories
        ORDER BY id DESC
        LIMIT ?
        """, (limit,)).fetchall()

        return [dict(r) for r in rows]

    def count(self):
        return self.db.execute(
            "SELECT COUNT(*) FROM memories"
        ).fetchone()[0]

    def clear(self):
        self.db.execute("DELETE FROM memories")
        self.db.commit()

    def close(self):
        self.db.close()