"""SQLite-backed memory index with embedding cache and hybrid BM25+vector search."""

from __future__ import annotations

import hashlib
import math
import sqlite3
import struct
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from loguru import logger

from asedclaw.agent.token_utils import estimate_tokens


# ---------------------------------------------------------------------------
# Data classes
# ---------------------------------------------------------------------------

@dataclass
class MemoryChunk:
    path: str  # relative to workspace
    start_line: int
    end_line: int
    text: str
    hash: str  # SHA-256 of text content


@dataclass
class SearchResult:
    path: str
    start_line: int
    end_line: int
    snippet: str
    score: float
    source: str  # "vector", "bm25", or "hybrid"


# ---------------------------------------------------------------------------
# Index manager
# ---------------------------------------------------------------------------

class MemoryIndexManager:
    """Chunking, embedding, and hybrid search over memory markdown files.

    Storage: ``workspace/memory/.index.db`` (SQLite).
    """

    CHUNK_TOKENS = 400
    CHUNK_OVERLAP = 80
    VECTOR_WEIGHT = 0.7
    TEXT_WEIGHT = 0.3
    CANDIDATE_MULTIPLIER = 4
    MAX_SNIPPET_CHARS = 700

    def __init__(
        self,
        workspace: Path,
        embedding_api_key: str | None = None,
        embedding_api_base: str | None = None,
        embedding_model: str = "jina-embeddings-v4",
        embedding_dimensions: int = 1024,
        embedding_task: str = "retrieval",
    ):
        self._workspace = workspace
        self._memory_dir = workspace / "memory"
        self._memory_file = self._memory_dir / "MEMORY.md"
        self._db_path = self._memory_dir / ".index.db"

        self._embedding_api_key = embedding_api_key
        self._embedding_api_base = embedding_api_base
        self._embedding_model = embedding_model
        self._embedding_dimensions = embedding_dimensions
        self._embedding_task = embedding_task

        self._conn: sqlite3.Connection | None = None
        self._init_db()

    # ------------------------------------------------------------------
    # Database setup
    # ------------------------------------------------------------------

    def _get_conn(self) -> sqlite3.Connection:
        if self._conn is None:
            self._memory_dir.mkdir(parents=True, exist_ok=True)
            self._conn = sqlite3.connect(str(self._db_path))
            self._conn.execute("PRAGMA journal_mode=WAL")
        return self._conn

    def _init_db(self) -> None:
        conn = self._get_conn()
        conn.executescript("""
            CREATE TABLE IF NOT EXISTS chunks (
                id INTEGER PRIMARY KEY,
                path TEXT NOT NULL,
                start_line INTEGER NOT NULL,
                end_line INTEGER NOT NULL,
                text TEXT NOT NULL,
                hash TEXT NOT NULL UNIQUE
            );

            CREATE TABLE IF NOT EXISTS embeddings (
                chunk_hash TEXT PRIMARY KEY,
                vector BLOB NOT NULL,
                model TEXT NOT NULL,
                created_at REAL NOT NULL
            );

            CREATE TABLE IF NOT EXISTS file_hashes (
                path TEXT PRIMARY KEY,
                content_hash TEXT NOT NULL,
                indexed_at REAL NOT NULL
            );
        """)
        # Create FTS5 table if not exists (cannot use IF NOT EXISTS in CREATE VIRTUAL TABLE
        # for some SQLite versions, so guard with a try).
        try:
            conn.execute(
                "CREATE VIRTUAL TABLE chunks_fts USING fts5(text, content='chunks', content_rowid='id')"
            )
        except sqlite3.OperationalError:
            pass  # already exists
        conn.commit()

    def close(self) -> None:
        if self._conn:
            self._conn.close()
            self._conn = None

    # ------------------------------------------------------------------
    # Chunking
    # ------------------------------------------------------------------

    @staticmethod
    def chunk_markdown(
        content: str,
        path: str,
        tokens: int = 400,
        overlap: int = 80,
    ) -> list[MemoryChunk]:
        """Split markdown content into overlapping chunks by line groups."""
        lines = content.splitlines()
        if not lines:
            return []

        chunks: list[MemoryChunk] = []
        current_lines: list[str] = []
        current_tokens = 0
        start_line = 1

        for i, line in enumerate(lines, start=1):
            line_tokens = estimate_tokens(line)
            current_lines.append(line)
            current_tokens += line_tokens

            if current_tokens >= tokens:
                text = "\n".join(current_lines)
                h = hashlib.sha256(text.encode("utf-8")).hexdigest()
                chunks.append(MemoryChunk(
                    path=path, start_line=start_line, end_line=i, text=text, hash=h,
                ))
                # Overlap: keep the last N tokens worth of lines
                overlap_lines: list[str] = []
                overlap_tokens = 0
                for ol in reversed(current_lines):
                    olt = estimate_tokens(ol)
                    if overlap_tokens + olt > overlap:
                        break
                    overlap_lines.insert(0, ol)
                    overlap_tokens += olt
                current_lines = overlap_lines
                current_tokens = overlap_tokens
                start_line = i - len(overlap_lines) + 1

        # Remaining lines
        if current_lines:
            text = "\n".join(current_lines)
            h = hashlib.sha256(text.encode("utf-8")).hexdigest()
            chunks.append(MemoryChunk(
                path=path, start_line=start_line, end_line=len(lines), text=text, hash=h,
            ))

        return chunks

    # ------------------------------------------------------------------
    # File scanning
    # ------------------------------------------------------------------

    def _list_memory_files(self) -> list[Path]:
        files: list[Path] = []
        if self._memory_file.exists():
            files.append(self._memory_file)
        if self._memory_dir.exists():
            for fp in sorted(self._memory_dir.glob("*.md")):
                if fp.name == "MEMORY.md":
                    continue
                files.append(fp)
        return files

    def _file_content_hash(self, path: Path) -> str:
        content = path.read_bytes()
        return hashlib.sha256(content).hexdigest()

    def _is_file_stale(self, path: Path) -> bool:
        conn = self._get_conn()
        rel = self._rel_path(path)
        row = conn.execute(
            "SELECT content_hash FROM file_hashes WHERE path = ?", (rel,)
        ).fetchone()
        if row is None:
            return True
        return row[0] != self._file_content_hash(path)

    def _rel_path(self, path: Path) -> str:
        try:
            return str(path.relative_to(self._workspace))
        except ValueError:
            return str(path)

    # ------------------------------------------------------------------
    # Sync
    # ------------------------------------------------------------------

    async def sync(self, force: bool = False) -> None:
        """Re-index stale or new memory files."""
        conn = self._get_conn()

        # Check for embedding model change
        self._clear_stale_cache()

        files = self._list_memory_files()
        indexed_paths: set[str] = set()

        for fp in files:
            rel = self._rel_path(fp)
            indexed_paths.add(rel)
            if not force and not self._is_file_stale(fp):
                continue

            logger.debug("Memory index: re-indexing {}", rel)
            content = fp.read_text(encoding="utf-8")
            content_hash = self._file_content_hash(fp)

            # Remove old chunks for this file
            old_ids = [
                r[0] for r in conn.execute("SELECT id FROM chunks WHERE path = ?", (rel,)).fetchall()
            ]
            if old_ids:
                placeholders = ",".join("?" * len(old_ids))
                conn.execute(f"DELETE FROM chunks WHERE id IN ({placeholders})", old_ids)
                conn.execute(f"DELETE FROM chunks_fts WHERE rowid IN ({placeholders})", old_ids)

            # Insert new chunks
            new_chunks = self.chunk_markdown(content, rel)
            for chunk in new_chunks:
                conn.execute(
                    "INSERT OR IGNORE INTO chunks (path, start_line, end_line, text, hash) VALUES (?, ?, ?, ?, ?)",
                    (chunk.path, chunk.start_line, chunk.end_line, chunk.text, chunk.hash),
                )

            # Rebuild FTS for this file's chunks
            rows = conn.execute(
                "SELECT id, text FROM chunks WHERE path = ?", (rel,)
            ).fetchall()
            for row_id, text in rows:
                conn.execute(
                    "INSERT OR REPLACE INTO chunks_fts (rowid, text) VALUES (?, ?)",
                    (row_id, text),
                )

            # Update file hash
            conn.execute(
                "INSERT OR REPLACE INTO file_hashes (path, content_hash, indexed_at) VALUES (?, ?, ?)",
                (rel, content_hash, time.time()),
            )

            # Embed new chunks
            await self._embed_chunks(new_chunks)

        # Clean up chunks for deleted files
        existing = {r[0] for r in conn.execute("SELECT path FROM file_hashes").fetchall()}
        removed = existing - indexed_paths
        for rel in removed:
            conn.execute("DELETE FROM chunks WHERE path = ?", (rel,))
            conn.execute("DELETE FROM file_hashes WHERE path = ?", (rel,))

        conn.commit()
        logger.debug("Memory index sync done: {} files", len(files))

    # ------------------------------------------------------------------
    # Embedding
    # ------------------------------------------------------------------

    async def _get_embedding(self, text: str) -> list[float] | None:
        """Get embedding from cache or API. Returns None if unavailable."""
        if not self._embedding_api_key:
            return None

        text_hash = hashlib.sha256(text.encode("utf-8")).hexdigest()
        conn = self._get_conn()

        # Check cache
        row = conn.execute(
            "SELECT vector FROM embeddings WHERE chunk_hash = ? AND model = ?",
            (text_hash, self._embedding_model),
        ).fetchone()
        if row:
            return self._unpack_vector(row[0])

        # Call embedding API (supports extra fields like task/dimensions for jina etc.)
        try:
            import httpx

            payload: dict[str, Any] = {
                "input": text,
                "model": self._embedding_model,
            }
            if self._embedding_dimensions:
                payload["dimensions"] = self._embedding_dimensions
            if self._embedding_task:
                payload["task"] = self._embedding_task

            base = (self._embedding_api_base or "").rstrip("/")
            async with httpx.AsyncClient(timeout=60) as client:
                resp = await client.post(
                    f"{base}/v1/embeddings",
                    json=payload,
                    headers={"Content-Type": "application/json"},
                )
                resp.raise_for_status()

            vector = resp.json()["data"][0]["embedding"]

            # Cache
            conn.execute(
                "INSERT OR REPLACE INTO embeddings (chunk_hash, vector, model, created_at) VALUES (?, ?, ?, ?)",
                (text_hash, self._pack_vector(vector), self._embedding_model, time.time()),
            )
            conn.commit()
            return vector
        except Exception:
            logger.debug("Embedding API call failed for chunk", exc_info=True)
            return None

    async def _embed_chunks(self, chunks: list[MemoryChunk]) -> None:
        """Embed chunks that are not yet cached."""
        if not self._embedding_api_key:
            return
        for chunk in chunks:
            await self._get_embedding(chunk.text)

    @staticmethod
    def _pack_vector(vec: list[float]) -> bytes:
        return struct.pack(f"{len(vec)}f", *vec)

    @staticmethod
    def _unpack_vector(blob: bytes) -> list[float]:
        n = len(blob) // 4
        return list(struct.unpack(f"{n}f", blob))

    def _clear_stale_cache(self) -> None:
        """Delete cached embeddings from a different model."""
        conn = self._get_conn()
        deleted = conn.execute(
            "DELETE FROM embeddings WHERE model != ?", (self._embedding_model,)
        ).rowcount
        if deleted:
            conn.commit()
            logger.info("Cleared {} stale embeddings (model changed)", deleted)

    # ------------------------------------------------------------------
    # Search: Vector
    # ------------------------------------------------------------------

    async def search_vector(
        self, query: str, max_results: int = 20,
    ) -> list[tuple[float, dict[str, Any]]]:
        """Cosine similarity search against stored embeddings."""
        query_vec = await self._get_embedding(query)
        if query_vec is None:
            return []

        conn = self._get_conn()
        rows = conn.execute(
            "SELECT c.path, c.start_line, c.end_line, c.text, c.hash "
            "FROM chunks c"
        ).fetchall()

        scored: list[tuple[float, dict[str, Any]]] = []
        for path, start_line, end_line, text, chunk_hash in rows:
            row = conn.execute(
                "SELECT vector FROM embeddings WHERE chunk_hash = ? AND model = ?",
                (chunk_hash, self._embedding_model),
            ).fetchone()
            if not row:
                continue
            chunk_vec = self._unpack_vector(row[0])
            sim = self._cosine_similarity(query_vec, chunk_vec)
            scored.append((sim, {
                "path": path, "start_line": start_line, "end_line": end_line,
                "snippet": text[:self.MAX_SNIPPET_CHARS],
            }))

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

    @staticmethod
    def _cosine_similarity(a: list[float], b: list[float]) -> float:
        if len(a) != len(b):
            return 0.0
        dot = sum(x * y for x, y in zip(a, b))
        norm_a = math.sqrt(sum(x * x for x in a))
        norm_b = math.sqrt(sum(x * x for x in b))
        if norm_a == 0 or norm_b == 0:
            return 0.0
        return dot / (norm_a * norm_b)

    # ------------------------------------------------------------------
    # Search: BM25 via FTS5
    # ------------------------------------------------------------------

    def search_bm25(
        self, query: str, max_results: int = 20,
    ) -> list[tuple[float, dict[str, Any]]]:
        """BM25 keyword search using SQLite FTS5."""
        conn = self._get_conn()

        # FTS5 query: split terms and join with OR for broader matching
        terms = query.split()
        fts_query = " OR ".join(f'"{t}"' for t in terms if t.strip())
        if not fts_query:
            return []

        try:
            rows = conn.execute(
                "SELECT f.rank, c.path, c.start_line, c.end_line, c.text "
                "FROM chunks_fts f "
                "JOIN chunks c ON c.id = f.rowid "
                "WHERE chunks_fts MATCH ? "
                "ORDER BY f.rank "
                "LIMIT ?",
                (fts_query, max_results * self.CANDIDATE_MULTIPLIER),
            ).fetchall()
        except sqlite3.OperationalError:
            logger.debug("FTS5 query failed, falling back", exc_info=True)
            return []

        results: list[tuple[float, dict[str, Any]]] = []
        for i, (rank, path, start_line, end_line, text) in enumerate(rows):
            # FTS5 rank is negative (lower = better), convert to 0-1 score
            score = 1.0 / (1.0 + max(0, i))
            results.append((score, {
                "path": path, "start_line": start_line, "end_line": end_line,
                "snippet": text[:self.MAX_SNIPPET_CHARS],
            }))

        return results[:max_results]

    # ------------------------------------------------------------------
    # Search: Hybrid
    # ------------------------------------------------------------------

    async def search_hybrid(
        self, query: str, max_results: int = 10,
    ) -> list[SearchResult]:
        """Merge vector and BM25 results with weighted fusion."""
        candidate_count = max_results * self.CANDIDATE_MULTIPLIER

        vector_results = await self.search_vector(query, candidate_count)
        bm25_results = self.search_bm25(query, candidate_count)

        # If no vector results, fall back to BM25 only
        if not vector_results and not bm25_results:
            return []

        # Build score maps keyed by (path, start_line)
        merged: dict[tuple[str, int], dict[str, Any]] = {}

        for score, info in vector_results:
            key = (info["path"], info["start_line"])
            if key not in merged:
                merged[key] = {**info, "vector_score": score, "bm25_score": 0.0}
            else:
                merged[key]["vector_score"] = max(merged[key].get("vector_score", 0), score)

        for score, info in bm25_results:
            key = (info["path"], info["start_line"])
            if key not in merged:
                merged[key] = {**info, "vector_score": 0.0, "bm25_score": score}
            else:
                merged[key]["bm25_score"] = max(merged[key].get("bm25_score", 0), score)

        # Compute weighted final score
        results: list[SearchResult] = []
        for entry in merged.values():
            vs = entry.get("vector_score", 0.0)
            bs = entry.get("bm25_score", 0.0)
            final = self.VECTOR_WEIGHT * vs + self.TEXT_WEIGHT * bs

            source = "hybrid"
            if vs > 0 and bs == 0:
                source = "vector"
            elif bs > 0 and vs == 0:
                source = "bm25"

            results.append(SearchResult(
                path=entry["path"],
                start_line=entry["start_line"],
                end_line=entry["end_line"],
                snippet=entry["snippet"],
                score=final,
                source=source,
            ))

        results.sort(key=lambda r: r.score, reverse=True)
        return results[:max_results]
