from pathlib import Path
from typing import List, Optional
import structlog

logger = structlog.get_logger()


class VectorStore:
    """
    向量存储 — 使用 ChromaDB 实现语义检索
    用于查找历史任务、架构模式、代码片段
    """

    def __init__(self, persist_dir: str = "./chroma_db", collection: str = "engineering"):
        self.persist_dir = persist_dir
        self.collection_name = collection
        self.logger = structlog.get_logger().bind(component="vector_store")
        self._client = None
        self._collection = None
        self._setup()

    def _setup(self):
        try:
            import chromadb
            from chromadb.config import Settings
            self._client = chromadb.PersistentClient(
                path=self.persist_dir,
                settings=Settings(anonymized_telemetry=False),
            )
            self._collection = self._client.get_or_create_collection(
                name=self.collection_name,
                metadata={"hnsw:space": "cosine"},
            )
            self.logger.info("VectorStore initialized", collection=self.collection_name)
        except ImportError:
            self.logger.warning("chromadb not installed, vector store disabled")
        except Exception as e:
            self.logger.warning("VectorStore setup failed", error=str(e))

    @property
    def available(self) -> bool:
        return self._collection is not None

    def add_document(self, doc_id: str, text: str, metadata: Optional[dict] = None) -> bool:
        if not self.available:
            return False
        try:
            self._collection.upsert(
                ids=[doc_id],
                documents=[text],
                metadatas=[metadata or {}],
            )
            return True
        except Exception as e:
            self.logger.warning("Failed to add document", error=str(e))
            return False

    def search(self, query: str, n_results: int = 5) -> List[dict]:
        if not self.available:
            return []
        try:
            results = self._collection.query(
                query_texts=[query],
                n_results=n_results,
            )
            items = []
            for i, doc in enumerate(results["documents"][0]):
                items.append({
                    "id": results["ids"][0][i],
                    "text": doc,
                    "distance": results["distances"][0][i] if results.get("distances") else None,
                    "metadata": results["metadatas"][0][i] if results.get("metadatas") else {},
                })
            return items
        except Exception as e:
            self.logger.warning("Search failed", error=str(e))
            return []

    def add_task_memory(self, task_id: str, user_request: str, architecture: dict, status: str):
        import json
        text = f"Request: {user_request}\nArchitecture: {json.dumps(architecture, ensure_ascii=False)[:500]}\nStatus: {status}"
        self.add_document(
            doc_id=task_id,
            text=text,
            metadata={"task_id": task_id, "status": status, "request": user_request[:200]},
        )

    def find_similar_architectures(self, user_request: str, n: int = 3) -> List[dict]:
        return self.search(f"Engineering request: {user_request}", n_results=n)

    def delete_document(self, doc_id: str) -> bool:
        if not self.available:
            return False
        try:
            self._collection.delete(ids=[doc_id])
            return True
        except Exception:
            return False