"""
vectors.py — soul-mcp Vector Layer

Vectorization layer for soul-mcp.
Responsibilities:
- Call local Ollama bge-m3 to generate embeddings
- Store vectors in soul_vectors table
- Semantic similarity recall

Principle: transparent to db.py, model can be swapped at any time.
Model standard: bge-m3 (native Chinese, multilingual, by BAAI)
History: 2026-03-04 migrated from nomic-embed-text to bge-m3
"""

import math
import struct
import sqlite3
import urllib.request
import urllib.error
import json
from datetime import datetime
from pathlib import Path

DB_PATH   = Path(__file__).parent / "soul.db"
MODEL     = "bge-m3"
OLLAMA_URL = "http://localhost:11434/api/embeddings"


# ── Vector encoding/decoding ──────────────────────────────────────

def _to_blob(vec: list) -> bytes:
    return struct.pack(f"{len(vec)}f", *vec)

def _from_blob(blob: bytes) -> list:
    n = len(blob) // 4
    return list(struct.unpack(f"{n}f", blob))

def _cosine(a: list, b: list) -> float:
    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))
    return dot / (mag_a * mag_b) if mag_a and mag_b else 0.0


# ── Ollama invocation ─────────────────────────────────────────────

def _augment_text(text: str, chunk_type: str = "") -> str:
    """
    Prompt Augmentation: add semantic context prefix for short text (< 30 chars)
    to improve vector quality. Long text is returned as-is (already has sufficient semantic info).
    """
    text = text.strip()
    if len(text) >= 30:
        return text
    prefix_map = {
        "emotion":       "Emotion/psychological state: ",
        "decision":      "Life decision/choice: ",
        "reason":        "Decision reason/rationale: ",
        "context":       "Situation/background at the time: ",
        "note":          "Personal note/thought: ",
        "topic":         "Silent topic/avoided area: ",
        "observation":   "Self-observation: ",
        "conflict":      "Inner contradiction/conflict: ",
        "analysis":      "Deep analysis: ",
        "principle":     "Life belief/principle: ",
        "identity":      "Self-cognition/identity: ",
        "memory_time":     "Past time memory: ",
        "prospect_time":   "Future prospect: ",
        "anchor_summary":  "Long-form material summary: ",
        "satellite":       "Key semantic tag: ",
        "paragraph":       "Long-form paragraph content: ",
    }
    if chunk_type == "query":
        return f"Search related soul memories: {text}"
    prefix = prefix_map.get(chunk_type, "Soul entry: ")
    return f"{prefix}{text}"


def get_embedding(text: str, chunk_type: str = "") -> list | None:
    """Call local Ollama to get vector, returns None on failure. Short text is auto-augmented."""
    if not text or not text.strip():
        return None
    text = _augment_text(text, chunk_type)
    payload = json.dumps({"model": MODEL, "prompt": text}).encode()
    req = urllib.request.Request(
        OLLAMA_URL,
        data=payload,
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    try:
        with urllib.request.urlopen(req, timeout=10) as resp:
            data = json.loads(resp.read())
            return data.get("embedding")
    except Exception:
        return None


def ollama_alive() -> bool:
    """Check if Ollama service is online."""
    try:
        urllib.request.urlopen("http://localhost:11434", timeout=2)
        return True
    except Exception:
        return False


# ── Single record write ────────────────────────────────────────────

def store_vector(soul_id: int, table_name: str, text: str) -> bool:
    """Generate and store a single vector, returns True on success."""
    vec = get_embedding(text)
    if vec is None:
        return False
    conn = sqlite3.connect(DB_PATH)
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute("PRAGMA busy_timeout=5000")
    # Only store once per record
    exists = conn.execute(
        "SELECT id FROM soul_vectors WHERE soul_id=? AND table_name=?",
        (soul_id, table_name)
    ).fetchone()
    if not exists:
        conn.execute(
            "INSERT INTO soul_vectors (soul_id, table_name, model, vector, created_at)"
            " VALUES (?,?,?,?,?)",
            (soul_id, table_name, MODEL, _to_blob(vec),
             datetime.now().isoformat())
        )
        conn.commit()
    conn.close()
    return True


# ── Semantic recall ────────────────────────────────────────────────

def semantic_search(query: str, limit: int = 6) -> list[tuple[int, str, float]]:
    """
    Semantic similarity search.
    Returns [(soul_id, table_name, score), ...], sorted by similarity descending.
    """
    query_vec = get_embedding(query, chunk_type="query")
    if query_vec is None:
        return []

    conn = sqlite3.connect(DB_PATH)
    rows = conn.execute(
        "SELECT soul_id, table_name, vector FROM soul_vectors WHERE vector IS NOT NULL"
    ).fetchall()
    conn.close()

    if not rows:
        return []

    scored = []
    for soul_id, table_name, blob in rows:
        vec = _from_blob(blob)
        score = _cosine(query_vec, vec)
        scored.append((soul_id, table_name, score))

    scored.sort(key=lambda x: x[2], reverse=True)
    return scored[:limit]


# ── Chunk-level vectors ────────────────────────────────────────────

def store_chunk_vectors(soul_id: int, table_name: str, chunks: list) -> int:
    """
    Update vectors for newly written soul_chunks rows.
    chunks: [(chunk_type, content), ...]
    Returns the number of chunks successfully vectorized.
    """
    if not chunks:
        return 0
    conn = sqlite3.connect(DB_PATH)
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute("PRAGMA busy_timeout=5000")
    done = 0
    for ctype, content in chunks:
        vec = get_embedding(content, chunk_type=ctype)
        if vec is None:
            conn.execute(
                "UPDATE soul_chunks SET vector_status='failed'"
                " WHERE soul_id=? AND table_name=? AND chunk_type=? AND vector IS NULL",
                (soul_id, table_name, ctype)
            )
            continue
        conn.execute(
            "UPDATE soul_chunks SET vector=?, model=?, vector_status='done'"
            " WHERE soul_id=? AND table_name=? AND chunk_type=? AND vector IS NULL",
            (_to_blob(vec), MODEL, soul_id, table_name, ctype)
        )
        done += 1
    conn.commit()
    conn.close()
    return done


def search_chunks(query: str, chunk_type: str = None,
                  table_name: str = None, limit: int = 6) -> list:
    """
    Chunk-level semantic search.
    Returns [{soul_id, table_name, chunk_type, content, score}, ...], sorted by similarity descending.
    Can filter by chunk_type (emotion/decision/reason/context/note/
    topic/observation/conflict/analysis/principle/identity)
    """
    query_vec = get_embedding(query, chunk_type="query")
    if query_vec is None:
        return []

    conn = sqlite3.connect(DB_PATH)
    sql = ("SELECT soul_id, table_name, chunk_type, content, vector, source_weight"
           " FROM soul_chunks WHERE vector IS NOT NULL")
    params = []
    if chunk_type:
        sql += " AND chunk_type=?"
        params.append(chunk_type)
    if table_name:
        sql += " AND table_name=?"
        params.append(table_name)
    rows = conn.execute(sql, params).fetchall()
    conn.close()

    if not rows:
        return []

    scored = []
    for row in rows:
        soul_id, tname, ctype, content, blob = row[0], row[1], row[2], row[3], row[4]
        sw = row[5] if len(row) > 5 else 1.0
        score = _cosine(query_vec, _from_blob(blob))
        scored.append({
            "soul_id": soul_id,
            "table_name": tname,
            "chunk_type": ctype,
            "content": content,
            "score": score,
            "source_weight": sw if sw is not None else 1.0,
        })

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


# ── Full backfill ────────────────────────────────────────────────────

def backfill(fts_content_fn, soul_tables, verbose=True) -> dict:
    """
    Generate vectors for all historical soul records.
    fts_content_fn: reference to db._fts_content function
    Returns {done, skipped, failed}
    """
    if not ollama_alive():
        return {"done": 0, "skipped": 0, "failed": 0, "error": "Ollama is not running"}

    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row  # Consistent with db._fts_content's named access
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute("PRAGMA busy_timeout=5000")
    already = {(r[0], r[1]) for r in
               conn.execute("SELECT soul_id, table_name FROM soul_vectors").fetchall()}

    done = skipped = failed = 0
    for table in soul_tables:
        try:
            rows = conn.execute(f"SELECT * FROM {table}").fetchall()
        except Exception:
            continue
        for row in rows:
            rid = row[0]
            if (rid, table) in already:
                skipped += 1
                continue
            text = fts_content_fn(table, row)
            if not text.strip():
                skipped += 1
                continue
            vec = get_embedding(text)
            if vec is None:
                failed += 1
                continue
            conn.execute(
                "INSERT INTO soul_vectors (soul_id, table_name, model, vector, created_at)"
                " VALUES (?,?,?,?,?)",
                (rid, table, MODEL, _to_blob(vec), datetime.now().isoformat())
            )
            done += 1
            if verbose and done % 10 == 0:
                print(f"  Vectorization progress: {done} records", flush=True)

    conn.commit()
    conn.close()
    return {"done": done, "skipped": skipped, "failed": failed}


def backfill_chunks(extract_fn, soul_tables, verbose=True) -> dict:
    """
    Generate chunk vectors for all historical records.
    extract_fn: reference to db._extract_chunks function
    Returns {done, skipped, failed}
    """
    if not ollama_alive():
        return {"done": 0, "skipped": 0, "failed": 0, "error": "Ollama is not running"}

    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row  # Consistent with db._extract_chunks's named access
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute("PRAGMA busy_timeout=5000")
    done = skipped = failed = 0

    for table in soul_tables:
        try:
            rows = conn.execute(f"SELECT * FROM {table}").fetchall()
        except Exception:
            continue

        for row in rows:
            rid = row[0]
            chunks = extract_fn(table, row)

            for ctype, content in chunks:
                # Check if already exists (skip if vector present)
                existing = conn.execute(
                    "SELECT id, vector FROM soul_chunks"
                    " WHERE soul_id=? AND table_name=? AND chunk_type=?",
                    (rid, table, ctype)
                ).fetchone()

                if existing and existing[1]:
                    skipped += 1
                    continue

                vec = get_embedding(content, chunk_type=ctype)
                if vec is None:
                    failed += 1
                    continue

                if existing:
                    conn.execute(
                        "UPDATE soul_chunks SET vector=?, model=?, vector_status='done' WHERE id=?",
                        (_to_blob(vec), MODEL, existing[0])
                    )
                else:
                    conn.execute(
                        "INSERT INTO soul_chunks"
                        " (soul_id, table_name, chunk_type, content, model, vector, created_at, vector_status)"
                        " VALUES (?,?,?,?,?,?,?,'done')",
                        (rid, table, ctype, content, MODEL,
                         _to_blob(vec), datetime.now().isoformat())
                    )
                done += 1
                if verbose and done % 20 == 0:
                    print(f"  Chunk vectorization progress: {done} chunks", flush=True)

    conn.commit()
    conn.close()
    return {"done": done, "skipped": skipped, "failed": failed}
