"""RAG Engine: Retrieve from ChromaDB + Generate via Agentic LLM."""
import logging
from typing import Dict, Any, List
from datetime import datetime
from app.core.database import query_documents, save_chat
from app.agent_workflow import agent

logger = logging.getLogger("deltamind.rag")

class RAGEngine:
    def __init__(self):
        self.counter = 0

    def _session_id(self) -> str:
        self.counter += 1
        return f"s_{datetime.now().strftime('%Y%m%d%H%M%S')}_{self.counter}"

    def retrieve(self, query: str, n: int = 4) -> Dict:
        results = query_documents(query, n)
        docs = results.get("documents", [[]])[0]
        metas = results.get("metadatas", [[]])[0]
        dists = results.get("distances", [[]])[0]
        
        parts, citations = [], []
        for i, (doc, meta, dist) in enumerate(zip(docs, metas, dists)):
            rel = max(0, round(1 - dist, 2))
            src = meta.get("source", "unknown")
            parts.append(f"[Source {i+1}: {src} | Relevance: {rel}]\n{doc}")
            citations.append({"index": i+1, "source": src, "relevance": rel})
            
        return {"context": "\n\n---\n\n".join(parts), "citations": citations, "count": len(docs)}

    async def chat(self, query: str, session_id: str = None, history: List[Dict] = None) -> Dict:
        sid = session_id or self._session_id()
        
        # 1. Retrieve
        ret = self.retrieve(query)
        
        # 2. Agent Generation (with potential tool calling)
        result = await agent.run(query, context=ret["context"], history=history)
        
        resp = {
            "session_id": sid, "query": query, "response": result.get("content",""),
            "provider": result.get("provider",""), "model": result.get("model",""),
            "route": result.get("route",""), "elapsed": result.get("elapsed",0),
            "citations": ret["citations"], "docs_retrieved": ret["count"],
            "timestamp": datetime.now().isoformat()
        }
        
        # 3. Save to history
        save_chat(sid, "user", query, result.get("provider",""))
        save_chat(sid, "assistant", result.get("content",""), result.get("provider",""))
        return resp

rag_engine = RAGEngine()