import os
import glob
import chromadb
import urllib.request
import json
try:
    from pypdf import PdfReader
    HAS_PYPDF = True
except ImportError:
    HAS_PYPDF = False

BASE_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../.."))
DB_DIR = os.path.join(BASE_DIR, "data", "vector_db")
KB_DIR = os.path.join(BASE_DIR, "data", "knowledge_base")

class OllamaEmbedder:
    def __init__(self, model="llama3"):
        self.model = model
        self.url = "http://localhost:11434/api/embeddings"
        
    def __call__(self, input: list[str]) -> list[list[float]]:
        embeddings = []
        for text in input:
            try:
                payload = {"model": self.model, "prompt": text}
                req = urllib.request.Request(self.url, data=json.dumps(payload).encode('utf-8'), headers={'Content-Type': 'application/json'})
                with urllib.request.urlopen(req) as response:
                    res = json.loads(response.read().decode('utf-8'))
                    embeddings.append(res.get("embedding", []))
            except Exception as e:
                print(f"[RAG ENGINE] Embedding error: {e}")
                embeddings.append([0.0] * 4096)
        return embeddings

try:
    chroma_client = chromadb.PersistentClient(path=DB_DIR)
    collection = chroma_client.get_or_create_collection(name="enterprise_knowledge", embedding_function=OllamaEmbedder(model="llama3"))
except Exception as e:
    collection = None

def chunk_text(text: str, chunk_size=1000, overlap=200) -> list[str]:
    chunks, start = [], 0
    while start < len(text):
        end = start + chunk_size
        chunks.append(text[start:end])
        start += chunk_size - overlap
    return chunks

def ingest_knowledge_base() -> str:
    if not collection: return "Error: ChromaDB offline."
    files = glob.glob(os.path.join(KB_DIR, "**/*"), recursive=True)
    valid = [f for f in files if os.path.isfile(f) and f.endswith(('.txt', '.md', '.pdf'))]
    if not valid: return "Ingestion halted: No documents in data/knowledge_base."
        
    count = 0
    for fp in valid:
        filename = os.path.basename(fp)
        text = ""
        if fp.endswith(".pdf") and HAS_PYPDF:
            try:
                for page in PdfReader(fp).pages:
                    extracted = page.extract_text()
                    if extracted: text += extracted + "\n"
            except Exception: pass
        else:
            try:
                with open(fp, "r", encoding="utf-8") as f: text = f.read()
            except Exception: pass
                
        if text.strip():
            chunks = chunk_text(text)
            ids = [f"{filename}_{i}" for i in range(len(chunks))]
            collection.add(documents=chunks, metadatas=[{"source": filename}]*len(chunks), ids=ids)
            count += len(chunks)
    return f"Knowledge Base Ingestion Complete. {count} chunks embedded."

def search_knowledge_base(query: str, n_results=3) -> str:
    if not collection: return "Error: ChromaDB offline."
    if collection.count() == 0: return "Knowledge Base is empty."
    try:
        results = collection.query(query_texts=[query], n_results=n_results)
        if not results["documents"] or not results["documents"][0]: return "No relevant info found."
        context = "--- RETRIEVED ENTERPRISE KNOWLEDGE ---\n"
        for idx, doc in enumerate(results["documents"][0]):
            source = results["metadatas"][0][idx].get("source", "Unknown")
            context += f"\n[Source: {source}]\n{doc}\n"
        return context
    except Exception as e: return f"Query Error: {str(e)}"
