import os
import json
import numpy as np
import faiss
import re 
from datasets import load_dataset
from sentence_transformers import SentenceTransformer
from google import genai
from sklearn.metrics.pairwise import cosine_similarity
import logging

logging.basicConfig(level=logging.INFO)

# CONFIG
API_KEY = os.getenv("GEMINI_API_KEY")
MODEL_NAME = "gemini-2.5-flash-lite"
if not API_KEY:
    raise ValueError("Missing GEMINI_API_KEY")
TOP_K = 4
MAX_MEMORY_ENTRIES = 10  # limit memory size
MEMORY_EMBED_DIM = 384    # use smaller embeddings for memory (MiniLM-L6)

HF_TOKEN = os.getenv("HF_TOKEN")
if not HF_TOKEN:
    raise ValueError("HF_TOKEN not found in environment variables") 
client = genai.Client(api_key=API_KEY)
embedder = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2", device="cpu")
memory_embedder = embedder 


# CACHE FILES
INDEX_FILE = "faiss.index"
DOC_FILE = "documents.json"
EMBED_FILE = "embeddings.npy"

embeddings = None 
#MEMORY STORAGE
memory_store = {}        # raw memory items per session
memory_embeddings = {}   # semantic embeddings per session 

# LOAD KNOWLEDGE
def load_knowledge():
    """
    Load or cache telecom knowledge documents safely.
    Works correctly with HuggingFace streaming datasets.
    """
    # ---------- LOAD CACHE ----------
    if os.path.exists(DOC_FILE):
        logging.info("✅ Loading cached documents")
        with open(DOC_FILE, "r") as f:
            return json.load(f)

    # ---------- LOAD DATASET (STREAMING SAFE) ----------
    try:
        ds = load_dataset(
            "electricsheepafrica/nigerian-telecom-customer-support-ticket-records",
            split="train",
            streaming=True,
            token=HF_TOKEN
        )
    except Exception as e:
        logging.warning(f"Dataset load failed: {e}")
        return []

    logging.info("⚙️ Streaming dataset and building knowledge base...")

    docs = []

    # ---------- DETECT COLUMNS FROM FIRST ROW ----------
    iterator = iter(ds)
    try:
        first_row = next(iterator)
    except StopIteration:
        logging.warning("Dataset empty")
        return []

    available_cols = list(first_row.keys())

    def safe_get(row, key):
        return row.get(key, "unknown") 
    # process first row
    rows = [first_row]
    
    # continue streaming remaining rows
    MAX_DOCS = 400 # start small

    for i, row in enumerate(iterator):
        if i >= MAX_DOCS:
            break
        rows.append(row)
    # ---------- BUILD DOCUMENTS ----------
    for i, row in enumerate(rows):
        text = f"""
Telecommunications Support Knowledge

Problem Type: {safe_get(row, 'issue_type')}

Observed Context:
- Operator: {safe_get(row, 'operator')}
- Channel: {safe_get(row, 'channel')}
- Priority: {safe_get(row, 'priority')}

Likely Resolution Pattern:
Cases with similar attributes were resolved in
{safe_get(row, 'resolution_time_hours')} hours
with customer satisfaction score
{safe_get(row, 'customer_satisfaction')}.

Use this as statistical troubleshooting evidence.
""".strip()

        docs.append({
            "id": f"DOC_{i}",
            "text": text
        })

    # ---------- CACHE ----------
    with open(DOC_FILE, "w") as f:
        json.dump(docs, f)

    logging.info(f"✅ Cached {len(docs)} telecom documents")

    return docs 
def get_kb():
    global documents, index, embeddings

    if documents is None or index is None:
        documents = load_knowledge()
        index, embeddings = load_or_create_index(documents)

    return documents, index 
if __name__ == "__main__":
    logging.info("Initializing Knowledge Base...")
    get_kb() 
    
#TRAINING DATA COLLECTION
def collect_training_data(user_input, answer, evaluation):
    global index, embeddings, documents 
    
    # 1. Ensure everything is initialized
    if index is None or documents is None:
        get_kb()

    # 2. Check the evaluation threshold
    if (evaluation.get("score", 0) >= 9 and 
        evaluation.get("grounded", False) and 
        not evaluation.get("needs_improvement", True) and len(answer.get("evidence_used", [])) >= 2
    ):
        
        entry_text = user_input + "\n" + json.dumps(answer)
        new_emb = embedder.encode([entry_text], normalize_embeddings=True)
        
        entry = {
            "id": f"TRAIN_{len(documents)}",
            "text": entry_text,
            "metadata": {
                "score": evaluation["score"],
                "evidence": answer.get("evidence_used", [])
            }
        }

        # 3. Semantic Deduplication
        if memory_embeddings.get("global") is None and embeddings is not None:
            memory_embeddings["global"] = embeddings.copy()
        
        global_mem = memory_embeddings.get("global")
        
        # Only add if it's unique enough (similarity < 0.9)
        if global_mem is None or np.max(cosine_similarity(new_emb, global_mem)) < 0.9:
            logging.info("Adding high-quality unique sample to Knowledge Base")
            
            # Update Document List
            documents.append(entry) 
            
            # Update FAISS Index (float32 is required by FAISS)
            index.add(new_emb.astype("float32"))
            
            # Update Global Embeddings for future deduplication
            embeddings = np.vstack([embeddings, new_emb])
            memory_embeddings["global"] = np.vstack([global_mem, new_emb]) if global_mem is not None else new_emb
            
            # Persistent Storage
            faiss.write_index(index, INDEX_FILE)
            np.save(EMBED_FILE, embeddings)
            
            # Save raw JSON for record keeping
            os.makedirs("auto_dataset", exist_ok=True)
            file_name = f"auto_dataset/{len(os.listdir('auto_dataset'))}.json"
            with open(file_name, "w") as f:
                json.dump(entry, f)
                
def rebuild_index():
    global index
    dim = embeddings.shape[1]
    index = faiss.IndexHNSWFlat(dim, 32, faiss.METRIC_INNER_PRODUCT)
    index.hnsw.efConstruction = 200
    index.add(embeddings.astype("float32"))
    faiss.write_index(index, INDEX_FILE) 
if os.path.exists(INDEX_FILE):
    logging.info("FAST START: loading cached index")    
# FAISS INDEX LOADING / BUILDING
def load_or_create_index(documents):
    global embeddings
    if os.path.exists(INDEX_FILE) and os.path.exists(EMBED_FILE):
        logging.info("✅ Loading cached FAISS index")
        index = faiss.read_index(INDEX_FILE)
        embeddings = np.load(EMBED_FILE)
        return index, embeddings

    logging.info("⚙️ Building embeddings (first run only)...")
    texts = [d["text"] for d in documents]
    embeddings = embedder.encode(texts, normalize_embeddings=True, show_progress_bar=True)

    # HNSW for incremental updates & RAM efficiency
    dim = embeddings.shape[1]
    index = faiss.IndexHNSWFlat(dim, 32, faiss.METRIC_INNER_PRODUCT)
    index.hnsw.efConstruction = 200
    index.add(np.array(embeddings).astype('float32'))

    faiss.write_index(index, INDEX_FILE)
    np.save(EMBED_FILE, embeddings)
    logging.info("✅ FAISS index cached")
    return index, embeddings
#index, embeddings = load_or_create_index(documents)
documents = None
index = None


# MEMORY MANAGEMENT
def get_memory(session="default"):
    """Return memory for a session."""
    return memory_store.get(session, [])

def update_memory(session, item):
    """
    Add item to session memory.
    Maintains MAX_MEMORY_ENTRIES limit and updates embeddings.
    """
    mem = memory_store.setdefault(session, [])
    mem.append(item)

    # Truncate if memory exceeds max entries
    if len(mem) > MAX_MEMORY_ENTRIES:
        # summarize old memory
        summary_prompt = f"""
Summarize the following telecom conversation to preserve context. Keep only key facts in JSON format if possible:
{mem}
"""
        try:
            summary_resp = client.models.generate_content(
                model=MODEL_NAME,
                contents=summary_prompt
            )
            summary_text = summary_resp.text
        except Exception as e:
            logging.warning(f"Memory summarization failed: {e}")
            summary_text = str(mem[-1])  # fallback: last entry

        # Replace memory with summarized entry
        mem = [{"summary": summary_text}]
        memory_store[session] = mem
        memory_embeddings[session] = memory_embedder.encode(
            [summary_text],
            normalize_embeddings=True
        )
    # Update embeddings for semantic retrieval
    latest = mem[-1]

    text = str(
        latest.get("summary")
        or latest.get("assistant")
        or latest.get("user")
        or ""
    )
    new_emb = memory_embedder.encode([text], normalize_embeddings=True)

    if session not in memory_embeddings:
        memory_embeddings[session] = new_emb
    else:
        memory_embeddings[session] = np.vstack(
            [memory_embeddings[session], new_emb]
        )
    return
# RETRIEVAL FUNCTION
def retrieve(query, session="default"):
    documents, index = get_kb()
    """
    Returns top-K documents + top memory entries based on similarity
    """
    results = []

    # ----- RAG from FAISS -----
    if len(documents) > 0:
        q_emb = embedder.encode([query], normalize_embeddings=True).astype('float32')
        if index is None or len(documents) == 0:
            return []
        D, I = index.search(q_emb, min(TOP_K, len(documents)))

        for idx in I[0]:
            results.append(documents[idx])
    # ----- Memory retrieval -----  
    mem_embs = memory_embeddings.get(session)
    if mem_embs is not None and len(mem_embs) > 0:
        q_emb_mem = memory_embedder.encode([query], normalize_embeddings=True)
        sim_scores = cosine_similarity(q_emb_mem, mem_embs)[0]

        top_idxs = np.argsort(sim_scores)[-2:]  # top 2 memories

        for idx in reversed(top_idxs):
           if sim_scores[idx] > 0.3:
                mem_entry = get_memory(session)[idx]
                results.append({
                    "id": f"MEM_{idx}",
                    "text": mem_entry.get("summary")
                        or mem_entry.get("assistant")
                        or ""
                })
    return results

# SYSTEM PROMPT
SYSTEM_PROMPT = """
You are a Tier-1 Telecommunications Support Engineer.

You MUST ground every diagnosis in provided evidence.

RULES:
- Use ONLY given evidence documents
- Cite evidence IDs used
- If evidence insufficient → request human agent
- No hallucination

Return STRICT JSON:

{
 "problem_category":"",
 "diagnosis":"",
 "recommended_steps":[],
 "evidence_used":[],
 "confidence_score":0-1,
 "requires_human_agent":false
}

Also determine problem_category from: network, billing, sim, device, configuration.
"""

# ROBUST JSON GENERATION
def generate_json(prompt):
    try:
        resp = client.models.generate_content(
            model=MODEL_NAME,
            contents=prompt,
            config={"response_mime_type": "application/json"}
        )
        return json.loads(resp.text)
    except Exception as e:
        logging.warning(f"JSON generation failed: {e}")
        return {"error": "Invalid JSON"}

# STREAMING ANSWER 
def stream_answer(prompt):
    buffer = ""
    try:
        stream = client.models.generate_content_stream(
            model=MODEL_NAME,
            contents=prompt,
            config={"response_mime_type": "application/json"}
        )
        for chunk in stream:
            if chunk.text:
                buffer += chunk.text
                if len(buffer) > 512:  # yield micro-batches
                    yield buffer
                    buffer = ""
        if buffer:
            yield buffer
    except Exception as e:
        logging.warning(f"Streaming failed: {e}")
        yield '{"text": "Streaming failure"}'
def rule_based_validation(answer):
    if not answer.get("evidence_used"):
        return False
    if answer.get("confidence_score", 0) > 0.95:
        return False
    if len(answer.get("diagnosis", "")) < 10:
        return False
    return True 
# AI JUDGE
JUDGE_PROMPT = """
You are a telecom AI evaluator.

Check:
1. Diagnosis supported by cited evidence
2. Evidence IDs exist in context
3. No unsupported claims
4. Logical troubleshooting

Return JSON:

{
 "score":0,
 "hallucination_risk":"",
 "grounded":true,
 "needs_improvement":false,
 "reason":""
}
"""
REQUIRED_KEYS = {
    "problem_category",
    "diagnosis",
    "recommended_steps",
    "evidence_used",
    "confidence_score",
    "requires_human_agent"
}

def valid_schema(answer):
    return isinstance(answer, dict) and REQUIRED_KEYS.issubset(answer.keys())
def judge(query, context, answer):
    prompt = f"""
{JUDGE_PROMPT}

Query:
{query}

Evidence Context:
{json.dumps(context)}

Assistant Response:
{json.dumps(answer)}
"""
    return generate_json(prompt)

def validate_evidence(answer, context):
    context_ids = {d["id"] for d in context}
    used_ids = set(answer.get("evidence_used", []))

    if not used_ids.issubset(context_ids):
        return False

    return True 

# SELF-IMPROVING RETRY LOOP
def generate_with_retry(query, context, memory, category, max_retry=1):
    last_answer = None
    last_eval = None
    for _ in range(max_retry):
        prompt = f"""
{SYSTEM_PROMPT}

Memory:
{json.dumps(memory)}


Context:
{json.dumps([{'id': d['id'], 'text': d['text']} for d in context])}

User Issue:
{query}

User Issue Category:
{category}
"""
        answer = generate_json(prompt)
        last_answer = answer
    
        if not valid_schema(answer):
            continue
        if not rule_based_validation(answer):
            continue
        if not validate_evidence(answer, context):
            answer["requires_human_agent"] = True
            answer["confidence_score"] = 0.3
        
        last_eval = judge(query, context, answer)
        # Retry if low confidence or requires human
        if answer.get("confidence_score", 1) < 0.7 or answer.get("requires_human_agent", False):
            l2_prompt = f"""
{SYSTEM_PROMPT}

You are retrying because confidence was low.
Re-evaluate STRICTLY using evidence only.

Memory:
{json.dumps(memory)}

Context:
{json.dumps([{'id': d['id'], 'text': d['text']} for d in context])}

User Issue:
{query}
"""
            answer = generate_json(l2_prompt)
            last_answer = answer
            last_eval = judge(query, context, answer)

        # Success threshold
        if last_eval.get("score", 0) >= 6 and last_eval.get("grounded", True):
            return last_answer, last_eval

    return last_answer, last_eval  

# ANALYTICS LOGGING
analytics = []

def log_session(query, result, evaluation, category):
    analytics.append({
        "query": query,
        "category": category,
        "score": evaluation.get("score", 0),
        "hallucination": evaluation.get("hallucination_risk", "")
    })

# MAIN TELECOM AGENT (non-streaming)
# =====================================================
def telecom_agent(user_input, session="default"):
    memory = get_memory(session)
    category = "unknown"
    context = retrieve(user_input, session)
    if not context:
        return {
            "problem_category": "unknown",
            "diagnosis": "No supporting evidence found",
            "recommended_steps": [],
            "evidence_used": [],
            "confidence_score": 0.0,
            "requires_human_agent": True
        }, {"score": 0}, []

    answer, evaluation = generate_with_retry(user_input, context, memory, category)
    update_memory(session, {"user": user_input, "assistant": answer})
    log_session(user_input, answer, evaluation, category)
    collect_training_data(user_input, answer, evaluation)

    # Return latest 5 analytics for insight
    return answer, evaluation, analytics[-5:]

def safe_parse(text):
    match = re.search(r'\{.*?\}', text, re.DOTALL)
    if match:
        try:
            return json.loads(match.group())
        except:
            return {"error": "corrupt json"}
    return {"error": "no json found"} 

# STREAMING VERSION
def telecom_agent_stream(user_input, session="default"):
    memory = get_memory(session)
    category = "unknown"
    context = retrieve(user_input, session)
    if not context:
        yield {"type": "final", "answer": {
             "problem_category": "unknown",
             "diagnosis": "No supporting evidence found",
             "recommended_steps": [],
             "evidence_used": [],
             "confidence_score": 0.0,
             "requires_human_agent": True
        }}
        return

    prompt = f"""
{SYSTEM_PROMPT}

Memory:
{json.dumps(memory)}

Context:
{json.dumps([{'id': d['id'], 'text': d['text']} for d in context])}

User Issue:
{user_input}
"""

    final_text = ""
    for partial in stream_answer(prompt):
        final_text += partial
        yield {"text": partial}
    # Robust final parsing
    try:
        answer = safe_parse(final_text)
    except Exception as e:
        logging.warning(f"Final JSON parsing failed: {e}")
        answer = {"error": "Invalid JSON"}
    evaluation = judge(user_input, context, answer)
    update_memory(session, {
        "user": user_input,
        "assistant": json.dumps(answer)
    })
    log_session(user_input, answer, evaluation, category)
    collect_training_data(user_input, answer, evaluation)

    yield {
        "type": "final",
        "answer": answer,
        "evaluation": evaluation,
        "analytics": analytics[-5:]
    }