import sqlite3
import json
import uuid
import struct
import urllib.request
import math
import os
import datetime
import logging
from pathlib import Path

logger = logging.getLogger("brynq.vector_store")

try:
    import numpy as np
    import faiss
    from sentence_transformers import SentenceTransformer
    HAS_VECTORS = True
except ImportError:
    HAS_VECTORS = False
    logger.warning("Vector dependencies missing (numpy, faiss-cpu, sentence-transformers). Falling back to basic embeddings.")

class VectorStore:
    def __init__(self, db_path: Path):
        self.db_path = db_path
        self.db_path.parent.mkdir(parents=True, exist_ok=True)
        self.conn = sqlite3.connect(str(self.db_path), check_same_thread=False)
        self.conn.row_factory = sqlite3.Row
        self._ensure_schema()
        
        self.embedder = None
        self.index = None
        self.id_map = {} 
        self.reverse_id_map = {}
        
        global HAS_VECTORS
        if HAS_VECTORS:
            try:
                self.embedder = SentenceTransformer('all-MiniLM-L6-v2')
                # all-MiniLM-L6-v2 outputs 384-dimensional embeddings
                self.index = faiss.IndexFlatL2(384)
                self._load_faiss_index()
            except Exception as e:
                logger.error(f"Failed to initialize sentence_transformers: {e}")
                HAS_VECTORS = False
                
    def _load_faiss_index(self):
        """Rebuild the FAISS index from SQLite on startup."""
        if not self.embedder:
            return
            
        c = self.conn.cursor()
        c.execute('SELECT id, content FROM vectors')
        rows = c.fetchall()
        
        if not rows:
            return
            
        texts = [r['content'] for r in rows]
        ids = [r['id'] for r in rows]
        
        try:
            embeddings = self.embedder.encode(texts, convert_to_numpy=True)
            self.index.add(embeddings)
            
            for i, doc_id in enumerate(ids):
                self.id_map[i] = doc_id
                self.reverse_id_map[doc_id] = i
        except Exception as e:
            logger.error(f"Failed to rebuild FAISS index: {e}")

    def _ensure_schema(self) -> None:
        c = self.conn.cursor()
        c.execute('''
            CREATE TABLE IF NOT EXISTS vectors (
                id TEXT PRIMARY KEY,
                content TEXT NOT NULL,
                embedding BLOB,
                metadata TEXT,
                namespace TEXT DEFAULT 'default',
                created_at TEXT,
                updated_at TEXT
            )
        ''')
        c.execute('CREATE INDEX IF NOT EXISTS idx_namespace ON vectors(namespace)')
        self.conn.commit()

    def _get_embedding(self, text: str) -> list[float]:
        try:
            req = urllib.request.Request(
                "http://localhost:11434/api/embeddings",
                data=json.dumps({"model": "nomic-embed-text", "prompt": text}).encode("utf-8"),
                headers={"Content-Type": "application/json"},
                method="POST"
            )
            with urllib.request.urlopen(req, timeout=2.0) as res:
                data = json.loads(res.read().decode("utf-8"))
                return data["embedding"]
        except Exception:
            from runtime.rag import _tokenize, _compute_tf
            import hashlib
            tf = _compute_tf(_tokenize(text))
            vec = [0.0] * 768
            for term, weight in tf.items():
                idx = int(hashlib.md5(term.encode('utf-8')).hexdigest(), 16) % 768
                vec[idx] += weight
            mag = math.sqrt(sum(v*v for v in vec))
            if mag > 0:
                vec = [v/mag for v in vec]
            return vec

    def _serialize_embedding(self, embedding: list[float]) -> bytes:
        return struct.pack('<' + 'f' * len(embedding), *embedding)

    def _deserialize_embedding(self, data: bytes) -> list[float]:
        if not data:
            return []
        count = len(data) // 4
        return list(struct.unpack('<' + 'f' * count, data))

    def add(self, content: str, metadata: dict | None = None, namespace: str = 'default', doc_id: str | None = None) -> str:
        doc_id = doc_id or uuid.uuid4().hex
        emb = self._get_embedding(content)
        blob = self._serialize_embedding(emb)
        now = datetime.datetime.now(datetime.timezone.utc).isoformat()
        
        c = self.conn.cursor()
        c.execute('''
            INSERT INTO vectors (id, content, embedding, metadata, namespace, created_at, updated_at)
            VALUES (?, ?, ?, ?, ?, ?, ?)
        ''', (doc_id, content, blob, json.dumps(metadata or {}), namespace, now, now))
        self.conn.commit()
        
        if self.embedder and self.index is not None:
            emb_faiss = self.embedder.encode([content], convert_to_numpy=True)
            curr_idx = self.index.ntotal
            self.index.add(emb_faiss)
            self.id_map[curr_idx] = doc_id
            self.reverse_id_map[doc_id] = curr_idx
            
        return doc_id

    def _cosine_similarity(self, a: list[float], b: list[float]) -> float:
        if not a or not b:
            return 0.0
        dot = sum(x * y for x, y in zip(a, b))
        mag_a = math.sqrt(sum(x * x for x in a))
        mag_b = math.sqrt(sum(x * x for x in b))
        if mag_a == 0 or mag_b == 0:
            return 0.0
        return dot / (mag_a * mag_b)

    def search(self, query: str, top_k: int = 5, namespace: str | None = None, min_score: float = 0.0) -> list[dict]:
        if self.embedder and self.index is not None and self.index.ntotal > 0:
            try:
                emb = self.embedder.encode([query], convert_to_numpy=True)
                search_k = min(top_k * 5, self.index.ntotal)
                distances, indices = self.index.search(emb, search_k)
                
                results = []
                c = self.conn.cursor()
                
                for i, idx in enumerate(indices[0]):
                    if idx == -1:
                        continue
                        
                    doc_id = self.id_map.get(idx)
                    if not doc_id:
                        continue
                        
                    c.execute('SELECT * FROM vectors WHERE id = ?', (doc_id,))
                    row = c.fetchone()
                    if not row:
                        continue
                        
                    if namespace and row['namespace'] != namespace:
                        continue
                        
                    d = dict(row)
                    d['metadata'] = json.loads(d['metadata']) if d['metadata'] else {}
                    # Invert distance for similarity score
                    score = float(1.0 / (1.0 + distances[0][i]))
                    if score >= min_score:
                        d['score'] = score
                        results.append(d)
                    
                    if len(results) >= top_k:
                        break
                        
                return results
            except Exception as e:
                logger.error(f"Vector search failed: {e}")
                # Fallback to local
                pass
                
        # Fallback to OLLAMA or TF-IDF embedding search
        q_emb = self._get_embedding(query)
        c = self.conn.cursor()
        if namespace:
            c.execute('SELECT * FROM vectors WHERE namespace = ?', (namespace,))
        else:
            c.execute('SELECT * FROM vectors')
            
        results = []
        for row in c.fetchall():
            emb = self._deserialize_embedding(row['embedding'])
            score = self._cosine_similarity(q_emb, emb)
            if score >= min_score:
                results.append({
                    'id': row['id'],
                    'content': row['content'],
                    'score': score,
                    'metadata': json.loads(row['metadata'] or '{}'),
                    'namespace': row['namespace'],
                    'created_at': row['created_at']
                })
                
        results.sort(key=lambda x: x['score'], reverse=True)
        return results[:top_k]

    def delete(self, doc_id: str) -> bool:
        c = self.conn.cursor()
        c.execute('DELETE FROM vectors WHERE id = ?', (doc_id,))
        deleted = c.rowcount > 0
        self.conn.commit()
        
        if deleted and self.embedder:
            # Rebuild FAISS index
            self.index.reset()
            self.id_map.clear()
            self.reverse_id_map.clear()
            self._load_faiss_index()
            
        return deleted

    def update(self, doc_id: str, content: str | None = None, metadata: dict | None = None) -> bool:
        c = self.conn.cursor()
        c.execute('SELECT * FROM vectors WHERE id = ?', (doc_id,))
        row = c.fetchone()
        if not row:
            return False
            
        new_content = content if content is not None else row['content']
        new_meta = metadata if metadata is not None else json.loads(row['metadata'])
        
        blob = row['embedding']
        if content is not None:
            blob = self._serialize_embedding(self._get_embedding(content))
            
        now = datetime.datetime.now(datetime.timezone.utc).isoformat()
        c.execute('''
            UPDATE vectors 
            SET content = ?, metadata = ?, embedding = ?, updated_at = ?
            WHERE id = ?
        ''', (new_content, json.dumps(new_meta), blob, now, doc_id))
        
        updated = c.rowcount > 0
        self.conn.commit()
        
        if updated and content is not None and self.embedder:
            # Updating FAISS index requires full rebuild since it's FlatL2
            self.index.reset()
            self.id_map.clear()
            self.reverse_id_map.clear()
            self._load_faiss_index()
            
        return True

    def get(self, doc_id: str) -> dict | None:
        c = self.conn.cursor()
        c.execute('SELECT * FROM vectors WHERE id = ?', (doc_id,))
        row = c.fetchone()
        if not row:
            return None
        return {
            'id': row['id'],
            'content': row['content'],
            'metadata': json.loads(row['metadata'] or '{}'),
            'namespace': row['namespace'],
            'created_at': row['created_at'],
            'updated_at': row['updated_at']
        }

    def list_namespaces(self) -> list[str]:
        c = self.conn.cursor()
        c.execute('SELECT DISTINCT namespace FROM vectors')
        return sorted([row['namespace'] for row in c.fetchall()])

    def count(self, namespace: str | None = None) -> int:
        c = self.conn.cursor()
        if namespace:
            c.execute('SELECT COUNT(*) FROM vectors WHERE namespace = ?', (namespace,))
        else:
            c.execute('SELECT COUNT(*) FROM vectors')
        row = c.fetchone()
        return row[0] if row else 0

    def batch_add(self, items: list[dict]) -> list[str]:
        now = datetime.datetime.now(datetime.timezone.utc).isoformat()
        rows = []
        ids = []
        for item in items:
            doc_id = item.get('id') or uuid.uuid4().hex
            content = item['content']
            emb = self._get_embedding(content)
            blob = self._serialize_embedding(emb)
            meta = json.dumps(item.get('metadata', {}))
            ns = item.get('namespace', 'default')
            rows.append((doc_id, content, blob, meta, ns, now, now))
            ids.append(doc_id)
            
        c = self.conn.cursor()
        c.executemany('''
            INSERT INTO vectors (id, content, embedding, metadata, namespace, created_at, updated_at)
            VALUES (?, ?, ?, ?, ?, ?, ?)
        ''', rows)
        self.conn.commit()
        return ids

    def batch_delete(self, doc_ids: list[str]) -> int:
        if not doc_ids:
            return 0
        c = self.conn.cursor()
        c.execute(f"DELETE FROM vectors WHERE id IN ({','.join('?'*len(doc_ids))})", doc_ids)
        self.conn.commit()
        return c.rowcount

    def search_by_metadata(self, filters: dict, namespace: str | None = None, limit: int = 20) -> list[dict]:
        c = self.conn.cursor()
        if namespace:
            c.execute('SELECT * FROM vectors WHERE namespace = ?', (namespace,))
        else:
            c.execute('SELECT * FROM vectors')
            
        results = []
        for row in c.fetchall():
            meta = json.loads(row['metadata'] or '{}')
            match = True
            for k, v in filters.items():
                if meta.get(k) != v:
                    match = False
                    break
            if match:
                results.append({
                    'id': row['id'],
                    'content': row['content'],
                    'metadata': meta,
                    'namespace': row['namespace'],
                    'created_at': row['created_at']
                })
                if len(results) >= limit:
                    break
        return results

    def clear_namespace(self, namespace: str) -> int:
        c = self.conn.cursor()
        c.execute('DELETE FROM vectors WHERE namespace = ?', (namespace,))
        self.conn.commit()
        return c.rowcount

    def vacuum(self) -> None:
        self.conn.execute('VACUUM')

    def export_namespace(self, namespace: str) -> list[dict]:
        c = self.conn.cursor()
        c.execute('SELECT id, content, metadata, created_at FROM vectors WHERE namespace = ?', (namespace,))
        return [{
            'id': row['id'],
            'content': row['content'],
            'metadata': json.loads(row['metadata'] or '{}'),
            'created_at': row['created_at']
        } for row in c.fetchall()]

    def import_namespace(self, namespace: str, data: list[dict]) -> int:
        for item in data:
            item['namespace'] = namespace
        self.batch_add(data)
        return len(data)

    def stats(self) -> dict:
        c = self.conn.cursor()
        c.execute('SELECT MIN(created_at), MAX(created_at) FROM vectors')
        oldest, newest = c.fetchone()
        return {
            'total_vectors': self.count(),
            'namespaces': self.list_namespaces(),
            'db_size_bytes': os.path.getsize(self.db_path) if self.db_path.exists() else 0,
            'oldest_entry': oldest,
            'newest_entry': newest
        }

    def prune_old(self, days: int = 90, namespace: str | None = None) -> int:
        cutoff = (datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(days=days)).isoformat()
        c = self.conn.cursor()
        if namespace:
            c.execute('DELETE FROM vectors WHERE namespace = ? AND created_at < ?', (namespace, cutoff))
        else:
            c.execute('DELETE FROM vectors WHERE created_at < ?', (cutoff,))
        self.conn.commit()
        return c.rowcount

    def reindex(self, namespace: str | None = None) -> int:
        c = self.conn.cursor()
        if namespace:
            c.execute('SELECT id, content FROM vectors WHERE namespace = ?', (namespace,))
        else:
            c.execute('SELECT id, content FROM vectors')
        
        rows = c.fetchall()
        count = 0
        for row in rows:
            blob = self._serialize_embedding(self._get_embedding(row['content']))
            c.execute('UPDATE vectors SET embedding = ? WHERE id = ?', (blob, row['id']))
            count += 1
        self.conn.commit()
        return count

    def deduplicate(self, namespace: str | None = None, similarity_threshold: float = 0.98) -> int:
        c = self.conn.cursor()
        if namespace:
            c.execute('SELECT id, embedding, created_at FROM vectors WHERE namespace = ? ORDER BY created_at ASC', (namespace,))
        else:
            c.execute('SELECT id, embedding, created_at FROM vectors ORDER BY created_at ASC')
            
        rows = c.fetchall()
        to_delete = []
        kept = []
        for row in rows:
            emb = self._deserialize_embedding(row['embedding'])
            is_dup = False
            for k_emb in kept:
                if self._cosine_similarity(emb, k_emb) >= similarity_threshold:
                    is_dup = True
                    break
            if is_dup:
                to_delete.append(row['id'])
            else:
                kept.append(emb)
                
        if to_delete:
            self.batch_delete(to_delete)
        return len(to_delete)

    def hybrid_search(self, query: str, keyword_weight: float = 0.3, semantic_weight: float = 0.7, top_k: int = 5, namespace: str | None = None) -> list[dict]:
        from runtime.rag import _tokenize
        q_tokens = set(_tokenize(query))
        
        c = self.conn.cursor()
        if namespace:
            c.execute('SELECT * FROM vectors WHERE namespace = ?', (namespace,))
        else:
            c.execute('SELECT * FROM vectors')
            
        q_emb = self._get_embedding(query)
        scored = []
        for row in c.fetchall():
            emb = self._deserialize_embedding(row['embedding'])
            sem_score = self._cosine_similarity(q_emb, emb)
            
            content_tokens = set(_tokenize(row['content']))
            overlap = len(q_tokens & content_tokens)
            kw_score = overlap / len(q_tokens) if q_tokens else 0.0
            
            final_score = (sem_score * semantic_weight) + (kw_score * keyword_weight)
            
            scored.append({
                'id': row['id'],
                'content': row['content'],
                'score': final_score,
                'metadata': json.loads(row['metadata'] or '{}'),
                'namespace': row['namespace'],
                'created_at': row['created_at']
            })
            
        scored.sort(key=lambda x: x['score'], reverse=True)
        return scored[:top_k]

    def close(self) -> None:
        self.conn.close()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()
