"""
Vector store backend implementations (unified backends module).

Moved from evenage.core.vector.
"""
from __future__ import annotations

from contextlib import contextmanager
import json
from pathlib import Path
from typing import Generator

from sqlalchemy import create_engine, text
from sqlalchemy.engine import Engine


class LocalVectorStore:
    """
    Local filesystem-based vector store for dev.
    Stores vectors as JSON files. Not suitable for production scale.
    """

    def __init__(self, storage_path: str = "./.evenage/vectors"):
        self.storage_path = Path(storage_path)
        self.storage_path.mkdir(parents=True, exist_ok=True)

    def store(self, key: str, vector: list[float], metadata: dict) -> None:
        data = {"key": key, "vector": vector, "metadata": metadata}
        file_path = self.storage_path / f"{self._sanitize_key(key)}.json"
        with open(file_path, "w") as f:
            json.dump(data, f)

    def search(self, vector: list[float], k: int = 5) -> list[dict]:
        results = []
        for file_path in self.storage_path.glob("*.json"):
            try:
                with open(file_path) as f:
                    data = json.load(f)
                stored_vector = data["vector"]
                similarity = self._cosine_similarity(vector, stored_vector)
                results.append({
                    "key": data["key"],
                    "metadata": data["metadata"],
                    "similarity": similarity,
                })
            except Exception:
                continue
        results.sort(key=lambda x: x["similarity"], reverse=True)
        return results[:k]

    def delete(self, key: str) -> None:
        fp = self.storage_path / f"{self._sanitize_key(key)}.json"
        if fp.exists():
            fp.unlink()

    def _sanitize_key(self, key: str) -> str:
        return key.replace("/", "_").replace("\\", "_").replace(":", "_")

    def _cosine_similarity(self, 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))
        ma = sum(x * x for x in a) ** 0.5
        mb = sum(y * y for y in b) ** 0.5
        if ma == 0 or mb == 0:
            return 0.0
        return dot / (ma * mb)


class PostgresVectorStore:
    """
    Postgres vector store using pgvector extension.
    Requires the database to support the 'vector' extension.
    """

    def __init__(self, database_url: str, dimension: int = 1536, table: str = "vectors"):
        self.database_url = database_url
        self.dimension = dimension
        self.table = table
        self._engine: Engine | None = None
        self._ensure_schema()

    def _get_engine(self) -> Engine:
        if self._engine is None:
            self._engine = create_engine(self.database_url, echo=False)
        return self._engine

    @contextmanager
    def _conn(self) -> Generator:
        engine = self._get_engine()
        with engine.begin() as conn:
            yield conn

    def _ensure_schema(self) -> None:
        try:
            with self._conn() as conn:
                conn.execute(text("CREATE EXTENSION IF NOT EXISTS vector"))
                conn.execute(text(
                    f"""
                    CREATE TABLE IF NOT EXISTS {self.table} (
                        key TEXT PRIMARY KEY,
                        embedding VECTOR({self.dimension}),
                        metadata JSONB
                    )
                    """
                ))
                conn.execute(text(
                    f"CREATE INDEX IF NOT EXISTS idx_{self.table}_embedding ON {self.table} USING ivfflat (embedding vector_cosine_ops) WITH (lists = 100)"
                ))
        except Exception:
            pass

    def store(self, key: str, vector: list[float], metadata: dict) -> None:
        if not isinstance(vector, list) or not vector:
            raise ValueError("vector must be a non-empty list[float]")
        vec_literal = "[" + ",".join(str(float(x)) for x in vector) + "]"
        with self._conn() as conn:
            conn.execute(
                text(
                    f"""
                    INSERT INTO {self.table} (key, embedding, metadata)
                    VALUES (:key, (:embedding)::vector, CAST(:metadata AS JSONB))
                    ON CONFLICT (key) DO UPDATE SET
                        embedding = EXCLUDED.embedding,
                        metadata = EXCLUDED.metadata
                    """
                ),
                {"key": key, "embedding": vec_literal, "metadata": json.dumps(metadata or {})},
            )

    def search(self, vector: list[float], k: int = 5) -> list[dict]:
        if not isinstance(vector, list) or not vector:
            raise ValueError("vector must be a non-empty list[float]")
        vec_literal = "[" + ",".join(str(float(x)) for x in vector) + "]"
        with self._conn() as conn:
            rows = conn.execute(
                text(
                    f"""
                    SELECT key, metadata,
                           1 - (embedding <=> (:qvec)::vector) AS similarity
                    FROM {self.table}
                    ORDER BY embedding <=> (:qvec)::vector
                    LIMIT :k
                    """
                ),
                {"qvec": vec_literal, "k": k},
            ).mappings().all()
        results: list[dict] = []
        for r in rows:
            try:
                meta = r["metadata"] if isinstance(r["metadata"], dict) else json.loads(r["metadata"]) if r["metadata"] else {}
            except Exception:
                meta = {}
            results.append({"key": r["key"], "metadata": meta, "similarity": float(r["similarity"])})
        return results

    def delete(self, key: str) -> None:
        with self._conn() as conn:
            conn.execute(text(f"DELETE FROM {self.table} WHERE key = :key"), {"key": key})
