app/services/engine.py

from __future__ import annotations

import time
import uuid
from typing import Any, Optional, Union

import numpy as np
from fastapi import HTTPException
from sqlalchemy.orm import Session

from app.models import Collection, Point, Snapshot
from app.schemas import (
    CollectionInfo,
    CountRequest,
    CreateCollectionRequest,
    DeletePointsRequest,
    Distance,
    PointStruct,
    RecommendRequest,
    Record,
    ScrollRequest,
    ScoredPoint,
    SearchRequest,
    UpdateResult,
    UpsertPointsRequest,
)
from app.services.embedding import bytes_to_vector, vector_to_bytes
from app.services.vector_math import filter_payload_fields, match_payload, score_vectors


def _now_result(result: Any, started: float) -> dict[str, Any]:
    return {"result": result, "status": "ok", "time": round(time.perf_counter() - started, 6)}


def _extract_vector(vector: Union[list[float], dict[str, list[float]]]) -> list[float]:
    if isinstance(vector, dict):
        if "" in vector:
            return vector[""]
        if len(vector) == 1:
            return next(iter(vector.values()))
        raise HTTPException(status_code=400, detail="Named vectors must include a single vector or '' key")
    return vector


class VectorService:
    def list_collections(self, db: Session) -> dict[str, Any]:
        started = time.perf_counter()
        rows = db.query(Collection).order_by(Collection.name).all()
        return _now_result({"collections": [{"name": c.name} for c in rows]}, started)

    def create_collection(self, db: Session, name: str, body: CreateCollectionRequest) -> dict[str, Any]:
        started = time.perf_counter()
        existing = db.query(Collection).filter(Collection.name == name).first()
        if existing:
            raise HTTPException(status_code=409, detail=f"Collection `{name}` already exists")

        hnsw = body.hnsw_config.model_dump() if body.hnsw_config else {
            "m": 16,
            "ef_construct": 100,
            "full_scan_threshold": 10000,
            "max_indexing_threads": 0,
            "on_disk": False,
        }
        optimizers = body.optimizers_config.model_dump() if body.optimizers_config else {
            "deleted_threshold": 0.2,
            "vacuum_min_vector_number": 1000,
            "default_segment_number": 0,
            "indexing_threshold": 20000,
            "flush_interval_sec": 5,
            "max_optimization_threads": None,
        }
        config = {
            "params": {
                "vectors": body.vectors.model_dump(),
                "shard_number": 1,
                "replication_factor": 1,
                "write_consistency_factor": 1,
                "on_disk_payload": body.on_disk_payload,
            },
            "hnsw_config": hnsw,
            "optimizer_config": optimizers,
        }

        row = Collection(
            name=name,
            vector_size=body.vectors.size,
            distance=body.vectors.distance.value if isinstance(body.vectors.distance, Distance) else body.vectors.distance,
            on_disk=body.vectors.on_disk,
            config=config,
            payload_schema={},
        )
        db.add(row)
        db.commit()
        return _now_result(True, started)

    def get_collection(self, db: Session, name: str) -> dict[str, Any]:
        started = time.perf_counter()
        row = self._get_collection_or_404(db, name)
        info = CollectionInfo(
            status=row.status,
            optimizer_status=row.optimizer_status,
            vectors_count=row.points_count,
            indexed_vectors_count=row.indexed_vectors_count,
            points_count=row.points_count,
            segments_count=row.segments_count,
            config=row.config or {},
            payload_schema=row.payload_schema or {},
        )
        return _now_result(info.model_dump(), started)

    def delete_collection(self, db: Session, name: str) -> dict[str, Any]:
        started = time.perf_counter()
        row = self._get_collection_or_404(db, name)
        db.query(Point).filter(Point.collection_id == row.id).delete()
        db.delete(row)
        db.commit()
        return _now_result(True, started)

    def collection_exists(self, db: Session, name: str) -> dict[str, Any]:
        started = time.perf_counter()
        exists = db.query(Collection).filter(Collection.name == name).first() is not None
        return _now_result(exists, started)

    def upsert_points(self, db: Session, name: str, body: UpsertPointsRequest) -> dict[str, Any]:
        started = time.perf_counter()
        col = self._get_collection_or_404(db, name)

        for p in body.points:
            vec = _extract_vector(p.vector)
            if len(vec) != col.vector_size:
                raise HTTPException(
                    status_code=400,
                    detail=f"Vector size mismatch: expected {col.vector_size}, got {len(vec)}",
                )
            point_id = str(p.id)
            existing = (
                db.query(Point)
                .filter(Point.collection_id == col.id, Point.point_id == point_id)
                .first()
            )
            payload = p.payload or {}
            self._update_payload_schema(col, payload)
            if existing:
                existing.vector = vector_to_bytes(vec)
                existing.payload = payload
            else:
                db.add(
                    Point(
                        collection_id=col.id,
                        point_id=point_id,
                        vector=vector_to_bytes(vec),
                        payload=payload,
                    )
                )

        db.flush()
        col.points_count = db.query(Point).filter(Point.collection_id == col.id).count()
        col.indexed_vectors_count = col.points_count
        db.commit()
        return _now_result(UpdateResult().model_dump(), started)

    def get_point(self, db: Session, name: str, point_id: str, with_payload: bool = True, with_vector: bool = False) -> dict[str, Any]:
        started = time.perf_counter()
        col = self._get_collection_or_404(db, name)
        point = (
            db.query(Point)
            .filter(Point.collection_id == col.id, Point.point_id == str(point_id))
            .first()
        )
        if not point:
            raise HTTPException(status_code=404, detail=f"Point `{point_id}` not found")
        rec = Record(
            id=_cast_id(point.point_id),
            payload=point.payload if with_payload else None,
            vector=bytes_to_vector(point.vector).tolist() if with_vector else None,
        )
        return _now_result(rec.model_dump(), started)

    def delete_points(self, db: Session, name: str, body: DeletePointsRequest) -> dict[str, Any]:
        started = time.perf_counter()
        col = self._get_collection_or_404(db, name)
        q = db.query(Point).filter(Point.collection_id == col.id)

        if body.points is not None:
            ids = [str(x) for x in body.points]
            q.filter(Point.point_id.in_(ids)).delete(synchronize_session=False)
        elif body.filter is not None:
            rows = q.all()
            for row in rows:
                if match_payload(row.payload, body.filter):
                    db.delete(row)
        else:
            raise HTTPException(status_code=400, detail="Provide points or filter")

        db.flush()
        col.points_count = db.query(Point).filter(Point.collection_id == col.id).count()
        col.indexed_vectors_count = col.points_count
        db.commit()
        return _now_result(UpdateResult().model_dump(), started)

    def search(self, db: Session, name: str, body: SearchRequest) -> dict[str, Any]:
        started = time.perf_counter()
        col = self._get_collection_or_404(db, name)
        query = np.asarray(_extract_vector(body.vector), dtype=np.float32)
        if query.shape[0] != col.vector_size:
            raise HTTPException(status_code=400, detail="Query vector size mismatch")

        points = db.query(Point).filter(Point.collection_id == col.id).all()
        filtered = [p for p in points if match_payload(p.payload, body.filter)]
        if not filtered:
            return _now_result([], started)

        matrix = np.vstack([bytes_to_vector(p.vector) for p in filtered])
        scores = score_vectors(query, matrix, col.distance)
        order = np.argsort(-scores)

        results: list[ScoredPoint] = []
        for idx in order:
            score = float(scores[idx])
            if body.score_threshold is not None and score < body.score_threshold:
                continue
            p = filtered[int(idx)]
            results.append(
                ScoredPoint(
                    id=_cast_id(p.point_id),
                    score=score,
                    payload=filter_payload_fields(p.payload, body.with_payload),
                    vector=bytes_to_vector(p.vector).tolist() if body.with_vector else None,
                )
            )
            if len(results) >= body.offset + body.limit:
                break

        sliced = results[body.offset : body.offset + body.limit]
        return _now_result([r.model_dump() for r in sliced], started)

    def recommend(self, db: Session, name: str, body: RecommendRequest) -> dict[str, Any]:
        started = time.perf_counter()
        col = self._get_collection_or_404(db, name)
        if not body.positive:
            raise HTTPException(status_code=400, detail="At least one positive example is required")

        pos_vecs = []
        for pid in body.positive:
            point = db.query(Point).filter(Point.collection_id == col.id, Point.point_id == str(pid)).first()
            if not point:
                raise HTTPException(status_code=404, detail=f"Positive point `{pid}` not found")
            pos_vecs.append(bytes_to_vector(point.vector))

        query = np.mean(np.vstack(pos_vecs), axis=0)
        if body.negative:
            neg_vecs = []
            for pid in body.negative:
                point = db.query(Point).filter(Point.collection_id == col.id, Point.point_id == str(pid)).first()
                if point:
                    neg_vecs.append(bytes_to_vector(point.vector))
            if neg_vecs:
                query = query - np.mean(np.vstack(neg_vecs), axis=0)

        exclude = {str(x) for x in body.positive + body.negative}
        search_body = SearchRequest(
            vector=query.tolist(),
            filter=body.filter,
            limit=body.limit + len(exclude),
            with_payload=body.with_payload,
            with_vector=body.with_vector,
            score_threshold=body.score_threshold,
        )
        raw = self.search(db, name, search_body)
        filtered = [r for r in raw["result"] if str(r["id"]) not in exclude][: body.limit]
        return _now_result(filtered, started)

    def scroll(self, db: Session, name: str, body: ScrollRequest) -> dict[str, Any]:
        started = time.perf_counter()
        col = self._get_collection_or_404(db, name)
        q = db.query(Point).filter(Point.collection_id == col.id).order_by(Point.id.asc())
        rows = q.all()
        filtered = [p for p in rows if match_payload(p.payload, body.filter)]

        start_idx = 0
        if body.offset is not None:
            for i, p in enumerate(filtered):
                if p.point_id == str(body.offset):
                    start_idx = i + 1
                    break

        page = filtered[start_idx : start_idx + body.limit]
        next_offset = page[-1].point_id if len(page) == body.limit and start_idx + body.limit < len(filtered) else None
        points = [
            Record(
                id=_cast_id(p.point_id),
                payload=filter_payload_fields(p.payload, body.with_payload),
                vector=bytes_to_vector(p.vector).tolist() if body.with_vector else None,
            ).model_dump()
            for p in page
        ]
        return _now_result({"points": points, "next_page_offset": next_offset}, started)

    def count(self, db: Session, name: str, body: CountRequest) -> dict[str, Any]:
        started = time.perf_counter()
        col = self._get_collection_or_404(db, name)
        rows = db.query(Point).filter(Point.collection_id == col.id).all()
        count = sum(1 for p in rows if match_payload(p.payload, body.filter))
        return _now_result({"count": count}, started)

    def create_snapshot(self, db: Session, name: str) -> dict[str, Any]:
        started = time.perf_counter()
        col = self._get_collection_or_404(db, name)
        snap_name = f"{name}-{int(time.time())}-{uuid.uuid4().hex[:8]}.snapshot"
        size = col.points_count * col.vector_size * 4
        snap = Snapshot(collection_name=name, name=snap_name, size_bytes=size)
        db.add(snap)
        db.commit()
        return _now_result({"name": snap_name}, started)

    def list_snapshots(self, db: Session, name: str) -> dict[str, Any]:
        started = time.perf_counter()
        self._get_collection_or_404(db, name)
        rows = db.query(Snapshot).filter(Snapshot.collection_name == name).order_by(Snapshot.created_at.desc()).all()
        return _now_result(
            [
                {
                    "name": s.name,
                    "creation_time": s.created_at.isoformat() if s.created_at else None,
                    "size": s.size_bytes,
                }
                for s in rows
            ],
            started,
        )

    def delete_snapshot(self, db: Session, name: str, snapshot_name: str) -> dict[str, Any]:
        started = time.perf_counter()
        row = (
            db.query(Snapshot)
            .filter(Snapshot.collection_name == name, Snapshot.name == snapshot_name)
            .first()
        )
        if not row:
            raise HTTPException(status_code=404, detail="Snapshot not found")
        db.delete(row)
        db.commit()
        return _now_result(True, started)

    def _get_collection_or_404(self, db: Session, name: str) -> Collection:
        row = db.query(Collection).filter(Collection.name == name).first()
        if not row:
            raise HTTPException(status_code=404, detail=f"Collection `{name}` doesn't exist")
        return row

    def _update_payload_schema(self, col: Collection, payload: dict[str, Any]) -> None:
        schema = dict(col.payload_schema or {})
        for key, value in payload.items():
            if key in schema:
                continue
            if isinstance(value, bool):
                t = "bool"
            elif isinstance(value, int):
                t = "integer"
            elif isinstance(value, float):
                t = "float"
            elif isinstance(value, dict):
                t = "object"
            elif isinstance(value, list):
                t = "list"
            else:
                t = "keyword"
            schema[key] = {"data_type": t, "points": 0}
        col.payload_schema = schema


def _cast_id(point_id: str) -> Union[int, str]:
    if point_id.isdigit():
        return int(point_id)
    return point_id


vector_service = VectorService()