"""
ChromaDB Vector Store Service
Ported from askme_open_agent/chromadb.py

API key resolution: agent_id (from DB) > openai_api_key > settings.OPENAI_API_KEYs
"""
import json
import uuid
from typing import List, Dict, Any, Optional
import chromadb
from chromadb.config import Settings

from app.core.config import settings
from app.services.agent_config import get_agent_config
from app.services.llm_service import LLMService


def _resolve_openai_api_key(
    agent_id: Optional[str] = None,
    openai_api_key: Optional[str] = None,
) -> str:
    """Resolve API key: from get_agent_config(agent_id), then openai_api_key, then settings."""
    if agent_id:
        try:
            cfg = get_agent_config(agent_id=agent_id)
            if cfg:
                key = (cfg.get("api_key") or cfg.get("open_ai") or "").strip()
                if key:
                    return key
        except Exception:
            pass
    if openai_api_key and str(openai_api_key).strip():
        return str(openai_api_key).strip()
    return (settings.OPENAI_API_KEY or "").strip()


class ChromaDBVectorStore:
    """ChromaDB vector store for training data (DDL, documentation, question-SQL pairs)"""
    
    def __init__(
        self,
        path: str = "./data/chromadb",
        collection_name: str = "askme_training",
        agent_id: Optional[str] = None,
        openai_api_key: Optional[str] = None,
    ):
        self.path = path
        self.collection_name = collection_name
        self._agent_id = agent_id
        self._openai_api_key = openai_api_key
        self._client = None
        self._collection = None
    
    def _get_client(self):
        """Get or create ChromaDB client"""
        if self._client is None:
            self._client = chromadb.PersistentClient(
                path=self.path,
                settings=Settings(anonymized_telemetry=False)
            )
        return self._client
    
    def _get_collection(self):
        """Get or create collection"""
        if self._collection is None:
            client = self._get_client()
            self._collection = client.get_or_create_collection(
                name=self.collection_name,
                metadata={"hnsw:space": "cosine"}
            )
        return self._collection
    
    def generate_embedding(self, text: str) -> List[float]:
        """Generate embedding using LLMService. API key from agent config (DB), openai_api_key, or settings."""
        api_key = _resolve_openai_api_key(
            agent_id=self._agent_id,
            openai_api_key=self._openai_api_key,
        )
        if not api_key:
            raise ValueError("OpenAI API key is not set. Provide agent_id, openai_api_key, or set OPENAI_API_KEY in config.")
        return LLMService(api_key=api_key).create_embedding(text)
    
    def add_question_sql(self, question: str, sql: str, **kwargs) -> str:
        """
        Add a question-SQL pair to training data
        
        Args:
            question: The question
            sql: The corresponding SQL query
            
        Returns:
            str: The ID of the training data
        """
        collection = self._get_collection()
        doc_id = f"{uuid.uuid4().hex[:8]}-sql"
        
        content = f"Question: {question}\nSQL: {sql}"
        embedding = self.generate_embedding(content)
        
        collection.add(
            ids=[doc_id],
            embeddings=[embedding],
            documents=[content],
            metadatas=[{
                "type": "question_sql",
                "question": question,
                "sql": sql
            }]
        )
        
        return doc_id
    
    def add_ddl(self, ddl: str, **kwargs) -> str:
        """
        Add DDL statement to training data
        
        Args:
            ddl: The DDL statement
            
        Returns:
            str: The ID of the training data
        """
        collection = self._get_collection()
        doc_id = f"{uuid.uuid4().hex[:8]}-ddl"
        
        embedding = self.generate_embedding(ddl)
        
        collection.add(
            ids=[doc_id],
            embeddings=[embedding],
            documents=[ddl],
            metadatas=[{
                "type": "ddl",
                "content": ddl[:500]  # First 500 chars for metadata
            }]
        )
        
        return doc_id
    
    def add_documentation(self, documentation: str, **kwargs) -> str:
        """
        Add documentation to training data
        
        Args:
            documentation: The documentation text
            
        Returns:
            str: The ID of the training data
        """
        collection = self._get_collection()
        doc_id = f"{uuid.uuid4().hex[:8]}-doc"
        
        embedding = self.generate_embedding(documentation)
        
        collection.add(
            ids=[doc_id],
            embeddings=[embedding],
            documents=[documentation],
            metadatas=[{
                "type": "documentation",
                "content": documentation[:500]
            }]
        )
        
        return doc_id
    
    def get_similar_question_sql(self, question: str, n_results: int = 5, **kwargs) -> List[Dict]:
        """
        Get similar question-SQL pairs for a given question
        
        Args:
            question: The question to search for
            n_results: Number of results to return
            
        Returns:
            List of similar question-SQL pairs
        """
        collection = self._get_collection()
        
        if collection.count() == 0:
            return []
        
        embedding = self.generate_embedding(question)
        
        results = collection.query(
            query_embeddings=[embedding],
            n_results=n_results,
            where={"type": "question_sql"}
        )
        
        similar = []
        if results and results['metadatas']:
            for metadata in results['metadatas'][0]:
                if 'question' in metadata and 'sql' in metadata:
                    similar.append({
                        "question": metadata['question'],
                        "sql": metadata['sql']
                    })
        
        return similar
    
    def get_related_ddl(self, question: str, n_results: int = 5, **kwargs) -> List[str]:
        """
        Get related DDL statements for a given question
        
        Args:
            question: The question to search for
            n_results: Number of results to return
            
        Returns:
            List of related DDL statements
        """
        collection = self._get_collection()
        
        if collection.count() == 0:
            return []
        
        embedding = self.generate_embedding(question)
        
        results = collection.query(
            query_embeddings=[embedding],
            n_results=n_results,
            where={"type": "ddl"}
        )
        
        if results and results['documents']:
            return results['documents'][0]
        
        return []
    
    def get_related_documentation(self, question: str, n_results: int = 5, **kwargs) -> List[str]:
        """
        Get related documentation for a given question
        
        Args:
            question: The question to search for
            n_results: Number of results to return
            
        Returns:
            List of related documentation
        """
        collection = self._get_collection()
        
        if collection.count() == 0:
            return []
        
        embedding = self.generate_embedding(question)
        
        results = collection.query(
            query_embeddings=[embedding],
            n_results=n_results,
            where={"type": "documentation"}
        )
        
        if results and results['documents']:
            return results['documents'][0]
        
        return []
    
    def get_training_data(self, **kwargs) -> List[Dict]:
        """
        Get all training data
        
        Returns:
            List of all training data items
        """
        collection = self._get_collection()
        
        if collection.count() == 0:
            return []
        
        results = collection.get()
        
        data = []
        if results:
            for i, doc_id in enumerate(results['ids']):
                data.append({
                    "id": doc_id,
                    "document": results['documents'][i] if results['documents'] else "",
                    "metadata": results['metadatas'][i] if results['metadatas'] else {}
                })
        
        return data
    
    def remove_training_data(self, id: str, **kwargs) -> bool:
        """
        Remove training data by ID
        
        Args:
            id: The ID of the training data
            
        Returns:
            bool: True if successful
        """
        collection = self._get_collection()
        
        try:
            collection.delete(ids=[id])
            return True
        except Exception:
            return False
    
    def train(self, ddl: str = None, documentation: str = None, question: str = None, sql: str = None, **kwargs) -> str:
        """
        Add training data (convenience method)
        
        Args:
            ddl: DDL statement
            documentation: Documentation text
            question: Question (requires sql)
            sql: SQL query (requires question)
            
        Returns:
            str: The ID of the added training data
        """
        if ddl:
            return self.add_ddl(ddl)
        elif documentation:
            return self.add_documentation(documentation)
        elif question and sql:
            return self.add_question_sql(question, sql)
        else:
            raise ValueError("Must provide ddl, documentation, or (question and sql)")


# Singleton instance
_chromadb_store = None


def get_chromadb_store(
    path: str = "./data/chromadb",
    collection_name: str = "askme_training",
    agent_id: Optional[str] = None,
    openai_api_key: Optional[str] = None,
) -> ChromaDBVectorStore:
    """
    Get or create ChromaDB vector store instance.

    API key is resolved from: agent_id (via get_agent_config from DB) > openai_api_key > settings.OPENAI_API_KEY.
    Pass agent_id when creating from AskmeService context to use the key stored in the database.
    """
    global _chromadb_store
    if _chromadb_store is None:
        _chromadb_store = ChromaDBVectorStore(
            path=path,
            collection_name=collection_name,
            agent_id=agent_id,
            openai_api_key=openai_api_key,
        )
    return _chromadb_store
