import asyncio
import json
import time
import traceback
from config import chat_client, CHAT_MODEL, RERANK_K
from rag.classifier import classify_question
from rag.query_expander import expand_query, contextualize_query
from rag.retriever import retrieve_hybrid, expand_context
from rag.reranker import rerank
from rag.prompt import get_system_prompt, build_user_message
from rag.utils import validate_and_clean_citations, merge_overlapping_chunks
from collections import Counter
import asyncio

# Answer-shape thresholds — a proxy for "is this a single fact or a multi-step
# procedure spanning several manual passages", measured from how much context
# survived retrieval + expansion. Deliberately based on the actual retrieved
# evidence rather than the question's intent category alone, since a "Maintenance"
# or "Troubleshooting" question can still resolve to one short spec line.
SHORT_ANSWER_CHAR_THRESHOLD = 600
LONG_ANSWER_CHAR_THRESHOLD = 1500
LONG_ANSWER_CHUNK_THRESHOLD = 3


def _determine_answer_shape(expanded_chunks: list[dict], context_str: str) -> str | None:
    """
    Infers "short" (single fact/spec) or "long" (multi-step procedure) from the size of
    the retrieved context, so build_user_message can give the model an explicit
    formatting instruction instead of relying on it to re-infer this from
    SYSTEM_PROMPT's static rules on every request. Returns None for the ambiguous
    middle ground, leaving the model to use its own default judgment.
    """
    char_count = len(context_str)
    chunk_count = len(expanded_chunks)
    if chunk_count <= 1 and char_count < SHORT_ANSWER_CHAR_THRESHOLD:
        return "short"
    if chunk_count >= LONG_ANSWER_CHUNK_THRESHOLD or char_count > LONG_ANSWER_CHAR_THRESHOLD:
        return "long"
    return None

class RAGCache:
    def __init__(self, ttl=3600):
        self.ttl = ttl
        self.local_cache = {}

    def get(self, machine_id: str, chat_key: str, msg: str):
        # machine_id is part of the key (not just chat_key/msg) so a cached answer built
        # from one machine's retrieved context can never be served back to a different
        # machine's chat asking the same normalized question text.
        cache_key = f"rag:cache:{machine_id}:{chat_key}:{msg}"
        print(f"[RAG Cache] Debug get key: '{cache_key}'", flush=True)
        # 1. Try Redis Cache first
        try:
            from database import redis_client
            if redis_client:
                cached = redis_client.get(cache_key)
                if cached:
                    print(f"[RAG Cache] Redis Hit! for query '{msg}' under session '{chat_key}'.", flush=True)
                    return json.loads(cached)
        except Exception as e:
            print(f"[RAG Cache] Redis get failed: {e}", flush=True)

        # 2. Fallback to Local Memory Dict Cache
        local_key = (machine_id, chat_key, msg)
        if local_key in self.local_cache:
            entry = self.local_cache[local_key]
            time_elapsed = time.time() - entry["timestamp"]
            if time_elapsed < self.ttl:
                print(f"[RAG Cache] Local Memory Hit! for query '{msg}' (Age: {int(time_elapsed)}s).", flush=True)
                return entry
            else:
                self.local_cache.pop(local_key, None)
        return None

    def set(self, machine_id: str, chat_key: str, msg: str, content: str, sources: list):
        cache_key = f"rag:cache:{machine_id}:{chat_key}:{msg}"
        print(f"[RAG Cache] Debug set key: '{cache_key}'", flush=True)
        data = {
            "content": content,
            "sources": sources,
            "timestamp": time.time()
        }
        # 1. Try Redis first
        try:
            from database import redis_client
            if redis_client:
                redis_client.setex(cache_key, self.ttl, json.dumps(data, default=str))
                print(f"[RAG Cache] Saved to Redis for query '{msg}'.", flush=True)
                return
        except Exception as e:
            print(f"[RAG Cache] Redis set failed: {e}", flush=True)

        # 2. Fallback to Local Memory
        local_key = (machine_id, chat_key, msg)
        self.local_cache[local_key] = data
        print(f"[RAG Cache] Saved to Local Memory for query '{msg}'.", flush=True)

    def clear(self, machine_id: str = None):
        """Invalidates cached queries for a machine (on ingest/delete/clear), or the whole cache."""
        if machine_id:
            keys_to_del = [k for k in list(self.local_cache.keys()) if k[0] == machine_id]
            for k in keys_to_del:
                self.local_cache.pop(k, None)
            print(f"[RAG Cache] Cleared {len(keys_to_del)} local cached queries for machine '{machine_id}'.", flush=True)
        else:
            count = len(self.local_cache)
            self.local_cache.clear()
            print(f"[RAG Cache] Cleared all {count} local cached queries.", flush=True)

        try:
            from database import redis_client
            if redis_client:
                pattern = f"rag:cache:{machine_id}:*" if machine_id else "rag:cache:*"
                deleted = 0
                for key in redis_client.scan_iter(match=pattern):
                    redis_client.delete(key)
                    deleted += 1
                print(f"[RAG Cache] Cleared {deleted} Redis cached queries matching '{pattern}'.", flush=True)
        except Exception as e:
            print(f"[RAG Cache] Redis clear failed: {e}", flush=True)

# Instantiate the global RAG Cache manager with 1 Hour TTL
_query_cache = RAGCache(ttl=3600)

def clear_query_cache(machine_id: str = None):
    """
    Public helper to invalidate cached RAG responses for a machine.
    Call this whenever a machine's knowledge base changes — after ingestion completes,
    after a document is deleted, or after the whole knowledge base is cleared — so a
    cached answer can never outlive the documents it was built from.
    """
    _query_cache.clear(machine_id)

async def sse_event(data: dict) -> str:
    """
    Formats a dictionary into an SSE event string.
    """
    return f"data: {json.dumps(data, default=str)}\n\n"

async def agent_stream(message: str, machine_id: str, history: list[dict]):
    """
    Main RAG v2 streaming pipeline (Non-streaming tokens, returning direct done/error):
    1. Greeting / Identity Check: Bypass RAG for simple inputs.
    2. Chat-Based TTL Cache Check: Check if query exists in cache for the active chat session (120s limit).
    3. Query Contextualization: Resolve follow-up pronouns/ellipsis into a standalone question using history.
    4. Intent Classification + Query Expansion: Run concurrently against the resolved query.
    5. Hybrid Retrieval: Run dense MMR and local BM25.
    6. Reranking: Rank results with FlashRank and take the top 10.
    7. Context Expansion: Expand the top 10 to include +/- 2 neighbors.
    8. Answer Shape: Infer short/long from retrieved context size.
    9. Prompt Construction: Inject expanded context + shape hint.
    10. LLM Generation: Fetch full answer silently.
    11. Answer Validation: Clean hallucinated citations from the final output.
    """
    # Send silent SSE comment to prevent Hugging Face proxy timeouts (completely hidden from Postman/JSON reader)
    yield ": ping\n\n"
    await asyncio.sleep(0)

    global _query_cache

    # Identify the chat key (using the first user message in history, or the current message if history is empty)
    chat_key = history[0].get("content", "") if (history and len(history) > 0) else message
    clean_chat_key = "".join(c for c in chat_key.lower() if c.isalnum() or c.isspace()).strip()

    # Normalize current query for cache lookups (strips punctuation like '?' and extra whitespaces)
    clean_msg = "".join(c for c in message.lower() if c.isalnum() or c.isspace()).strip()
    
    # 1. Check Cache (Redis with memory fallback)
    cached_entry = _query_cache.get(machine_id, clean_chat_key, clean_msg)
    if cached_entry:
        yield await sse_event({"type": "done", "content": cached_entry["content"], "sources": cached_entry["sources"]})
        return

    print(f"[RAG v2 Query] Machine: {machine_id} | User Asked: \"{message}\"", flush=True)
    
    # Quick Local Check for Greetings & Identity
    greetings = {
        "hi", "hello", "hellos", "hey", "greetings", "good morning", "good afternoon", "good evening", 
        "howdy", "assalam o alaikum", "salam", "ola", "hy", "hii", "helo", "heyy", "yo", "yoo", "aoa", "slm", "ok", "thanks", "thankyou", "thank you", "nice", "good", "great", "fine", "cool", "perfect"
    }
    who_are_you = {"who are you", "what is this", "what do you do", "who is this", "what are you", "how are you"}
    
    if clean_msg in greetings or clean_msg.startswith(("hello ", "hi ", "hey ", "salam ", "hy ", "yo ", "hii ", "helo ")):
        reply = "Hello! I am your factory diagnostics assistant. Ask me anything about this machine's operations or maintenance, and please ensure the manual is uploaded for specific queries!"
        yield await sse_event({"type": "done", "content": reply, "sources": []})
        return
        
    if clean_msg in who_are_you:
        reply = "I am the factory floor diagnostics AI. I can parse machine manuals and answer your technical questions, guide you through checklists, or extract critical safety warnings!"
        yield await sse_event({"type": "done", "content": reply, "sources": []})
        return

    # 2. Query Contextualization: resolve follow-up references ("how often should I do
    # that") into a standalone question using recent history, BEFORE classification/
    # expansion/retrieval run — those all only ever saw the raw message otherwise, so a
    # follow-up's pronouns and ellipsis were invisible to search. No-ops (no LLM call)
    # when there's no history yet.
    resolved_query = await asyncio.to_thread(contextualize_query, message, history)
    if resolved_query != message:
        print(f"RAG v2: Contextualized query: \"{message}\" -> \"{resolved_query}\"", flush=True)

    # 3. Intent Classification + Query Expansion — run concurrently since neither
    # depends on the other's output, both against the resolved query.
    intent, expanded_queries = await asyncio.gather(
        asyncio.to_thread(classify_question, resolved_query),
        asyncio.to_thread(expand_query, resolved_query)
    )
    print(f"RAG v2: Intent classified as '{intent}'", flush=True)
    print(f"RAG v2: Expanded queries: {expanded_queries}", flush=True)

    # 4. Hybrid Retrieval
    try:
        merged_chunks = await asyncio.to_thread(retrieve_hybrid, expanded_queries, machine_id, intent)
    except Exception as e:
        print(f"Error in retrieve_hybrid: {e}")
        traceback.print_exc()
        merged_chunks = []

    if not merged_chunks:
        yield await sse_event({"type": "done", "content": "Could not find this in the uploaded documents.", "sources": []})
        return

    # 5. FlashRank Reranking (Top 40 -> Top 8)
    try:
        top_reranked = await asyncio.to_thread(rerank, resolved_query, merged_chunks, top_n=RERANK_K)
    except Exception as e:
        print(f"Error in reranking: {e}")
        top_reranked = merged_chunks[:RERANK_K]
 
        
    # 6. Context Expansion (Retrieve neighbors +/- 2 chunks)
    try:
        expanded_chunks = await asyncio.to_thread(expand_context, top_reranked, machine_id)
    except Exception as e:
        print(f"Error in context expansion: {e}")
        expanded_chunks = top_reranked

    # # Lock context to the top document to prevent cross-document answer mixing
    # top_filename = top_reranked[0]["filename"] if top_reranked else None
    # if top_filename:
    #     expanded_chunks = [c for c in expanded_chunks if c["filename"] == top_filename]
    #     print(f"RAG v2: Context locked to '{top_filename}' ({len(expanded_chunks)} chunks)", flush=True)
    
    # Yield retrieved chunks metadata for frontend citation panel (use expanded_chunks so all cited pages are present)
    retrieved_event_data = []
    seen_pages = set()
    for c in expanded_chunks:
        key = (c["filename"], c["page"])
        if key not in seen_pages:
            retrieved_event_data.append({
                "filename": c["filename"],
                "page": c["page"],
                "score": c.get("score", 0.0)
            })
            seen_pages.add(key)
        
    # 7. Prompt Construction — merge chunk-linked adjacent entries first, stripping the
    # chunker's own 120-char overlap so the LLM doesn't see the same row twice across
    # two separate [Source: ...] blocks (confirmed live: that duplicate line reads as a
    # natural stopping point and silently truncates real checklist rows that followed).
    expanded_chunks = merge_overlapping_chunks(expanded_chunks)

    context_str = "\n\n---\n\n".join(
        f"[Source: {c['filename']} | page {c['page']}]\n{c['text']}"
        for c in expanded_chunks
    )

    answer_shape = _determine_answer_shape(expanded_chunks, context_str)
    if answer_shape:
        print(f"RAG v2: Answer shape hint '{answer_shape}' ({len(expanded_chunks)} chunks, {len(context_str)} chars)", flush=True)

    rag_user_message = build_user_message(context_str, resolved_query, answer_shape=answer_shape)

    messages = [{"role": "system", "content": get_system_prompt()}]
    for h in history:
        messages.append(h)
    messages.append({"role": "user", "content": rag_user_message})

    # 8. Fetch LLM response silently (Non-streamed call for faster response)
    try:
        response = await asyncio.to_thread(
            chat_client.chat.completions.create,
            model=CHAT_MODEL,
            messages=messages,
            max_tokens=4096,
            temperature=0.0,
        )
        full_answer = response.choices[0].message.content or ""
                
        # 9. Answer Validation (Clean hallucinated citations from the final saved text)
        cleaned_answer, had_hallucination = validate_and_clean_citations(full_answer, expanded_chunks)
        if had_hallucination:
            print("RAG v2: Cleaned answer by removing hallucinated citations.")
            
        # Only extract unique sources that were ACTUALLY cited by the LLM (keeps done event payload tiny!)
        import re
        citation_matches = re.findall(
            r'\[Source:\s*([^\\|\]]+?)\s*\\?\|\s*(?:pages?|ch\.)?\s*([\d\s\-,]+)\s*\]',
            cleaned_answer,
            re.IGNORECASE
        )
        cited_keys = set()
        for fname, page_str in citation_matches:
            fname = fname.strip()
            page_num_match = re.search(r'\d+', page_str)
            if page_num_match:
                pnum = int(page_num_match.group(0))
                cited_keys.add((fname, pnum))
                
        filtered_sources = []
        seen_pages = set()
        for src in retrieved_event_data:
            key = (src["filename"], src["page"])
            if key in cited_keys and key not in seen_pages:
                filtered_sources.append(src)
                seen_pages.add(key)
                
        # Fallback: if no citations parsed but we have sources, just send unique top 2 to avoid empty panel
        if not filtered_sources and retrieved_event_data:
            for src in retrieved_event_data:
                key = (src["filename"], src["page"])
                if key not in seen_pages:
                    filtered_sources.append(src)
                    seen_pages.add(key)
                    if len(filtered_sources) >= 2:
                        break
            
        # Save validated response, parsed sources into cache (Redis / Local Memory) — but
        # never cache a "not found" answer, since that's exactly the case most likely to
        # change soon (a document finishing ingestion, or a new upload) and staying wrong
        # for the cache's full TTL would be worse than just re-running retrieval.
        is_not_found = "could not find" in cleaned_answer.lower()
        if not is_not_found:
            _query_cache.set(machine_id, clean_chat_key, clean_msg, cleaned_answer, filtered_sources)

        yield await sse_event({"type": "done", "content": cleaned_answer, "sources": filtered_sources})
        
    except Exception as e:
        print("\n" + "="*60)
        print("Chat System Error Encountered")
        traceback.print_exc()
        print("="*60 + "\n")
        err = str(e)
        if "rate_limit" in err.lower() or "429" in err:
            yield await sse_event({"type": "error", "content": "Rate limit reached. Please wait 1–2 minutes and try again."})
        else:
            yield await sse_event({"type": "error", "content": f"API Error: {err}"})

async def agent_chat(message: str, machine_id: str, history: list[dict]) -> dict:
    """
    Non-streaming implementation with full terminal diagnostics print statements.
    """
    print(f"\n" + "="*80, flush=True)
    print(f"[RAG Diagnostics] Query Received: \"{message}\" (Machine ID: {machine_id})", flush=True)
    print("="*80, flush=True)

    # 1. Quick Check for Greetings & Identity
    clean_msg = "".join(c for c in message.lower() if c.isalnum() or c.isspace()).strip()
    greetings = {
        "hi", "hello", "hellos", "hey", "greetings", "good morning", "good afternoon", "good evening", 
        "howdy", "assalam o alaikum", "salam", "ola", "hy", "hii", "helo", "heyy", "yo", "yoo", "aoa", "slm", "ok", "thanks", "thankyou", "thank you", "nice", "good", "great", "fine", "cool", "perfect"
    }
    who_are_you = {"who are you", "what is this", "what do you do", "who is this", "what are you", "how are you"}
    
    if clean_msg in greetings or clean_msg.startswith(("hello ", "hi ", "hey ", "salam ", "hy ", "yo ", "hii ", "helo ")):
        reply = "Hello! I am your factory diagnostics assistant. Ask me anything about this machine's operations or maintenance, and please ensure the manual is uploaded for specific queries!"
        print(f"[RAG Diagnostics] Auto-replying to greeting: \"{reply}\"", flush=True)
        return {"status": "done", "content": reply, "sources": [], "searches": []}
        
    if clean_msg in who_are_you:
        reply = "I am the factory floor diagnostics AI. I can parse machine manuals and answer your technical questions, guide you through checklists, or extract critical safety warnings!"
        print(f"[RAG Diagnostics] Auto-replying to identity query: \"{reply}\"", flush=True)
        return {"status": "done", "content": reply, "sources": [], "searches": []}

    # ── 1.1 Check Cache (Redis with memory fallback) ──
    chat_key = history[0].get("content", "") if (history and len(history) > 0) else message
    clean_chat_key = "".join(c for c in chat_key.lower() if c.isalnum() or c.isspace()).strip()
    
    cached_entry = _query_cache.get(machine_id, clean_chat_key, clean_msg)
    if cached_entry:
        print(f"[RAG Cache] Hit in agent_chat! Returning cached response for '{clean_msg}'.", flush=True)
        return {
            "status": "done",
            "content": cached_entry["content"],
            "sources": cached_entry["sources"],
            "searches": [{"query": message, "done": True, "count": len(cached_entry["sources"])}]
        }

    # 2. Query Contextualization (see agent_stream for why): resolve follow-up
    # references into a standalone question before classification/expansion/retrieval.
    resolved_query = await asyncio.to_thread(contextualize_query, message, history)
    if resolved_query != message:
        print(f"[RAG Diagnostics] Contextualized query: \"{message}\" -> \"{resolved_query}\"", flush=True)

    # 3. Intent Classification + Query Expansion — run concurrently, against the
    # resolved query.
    intent, expanded_queries = await asyncio.gather(
        asyncio.to_thread(classify_question, resolved_query),
        asyncio.to_thread(expand_query, resolved_query)
    )
    print(f"[RAG Diagnostics] Intent Category: '{intent}'", flush=True)
    print(f"[RAG Diagnostics] Expanded Queries: {expanded_queries}", flush=True)

    # 4. Hybrid Retrieval
    try:
        merged_chunks = await asyncio.to_thread(retrieve_hybrid, expanded_queries, machine_id, intent)
        print(f"[RAG Diagnostics] Hybrid retrieval found {len(merged_chunks)} raw chunks.", flush=True)
    except Exception as e:
        print(f"[RAG Diagnostics] ERROR in retrieve_hybrid: {e}", flush=True)
        merged_chunks = []

    if not merged_chunks:
        print("[RAG Diagnostics] No chunks found in database. Returning fallback response.", flush=True)
        return {
            "status": "done",
            "content": "Could not find this in the uploaded documents.",
            "sources": [],
            "searches": [{"query": message, "done": True, "count": 0}]
        }

    # 5. FlashRank Reranking
    try:
        top_reranked = await asyncio.to_thread(rerank, resolved_query, merged_chunks, top_n=RERANK_K)
        print(f"[RAG Diagnostics] Reranking selected top {len(top_reranked)} chunks.", flush=True)
    except Exception as e:
        print(f"[RAG Diagnostics] ERROR in reranking: {e}", flush=True)
        top_reranked = merged_chunks[:RERANK_K]
        
    # 6. Context Expansion
    try:
        expanded_chunks = await asyncio.to_thread(expand_context, top_reranked, machine_id)
        print(f"[RAG Diagnostics] Context expanded (+/-2 neighbors) to {len(expanded_chunks)} total chunks.", flush=True)
    except Exception as e:
        print(f"[RAG Diagnostics] ERROR in context expansion: {e}", flush=True)
        expanded_chunks = top_reranked

    # Format sources and searches
    sources = []
    valid_sources_log = []
    for c in expanded_chunks:
        sources.append({
            "filename": c["filename"],
            "page": c["page"],
            "score": c.get("score", 0.0)
        })
        valid_sources_log.append((c["filename"], c["page"]))
        
    print(f"[RAG Diagnostics] Valid sources in context: {valid_sources_log}", flush=True)
    searches = [{"query": message, "done": True, "count": len(sources)}]
    
    # 7. Prompt Construction — same overlap-deduplication as agent_stream (see there).
    expanded_chunks = merge_overlapping_chunks(expanded_chunks)

    context_str = "\n\n---\n\n".join(
        f"[Source: {c['filename']} | page {c['page']}]\n{c['text']}"
        for c in expanded_chunks
    )

    answer_shape = _determine_answer_shape(expanded_chunks, context_str)
    if answer_shape:
        print(f"[RAG Diagnostics] Answer shape hint '{answer_shape}' ({len(expanded_chunks)} chunks, {len(context_str)} chars)", flush=True)

    rag_user_message = build_user_message(context_str, resolved_query, answer_shape=answer_shape)

    messages = [{"role": "system", "content": get_system_prompt()}]
    for h in history:
        messages.append(h)
    messages.append({"role": "user", "content": rag_user_message})

    # 8. Call LLM (non-streamed)
    try:
        print(f"[RAG Diagnostics] Sending prompt to LLM ({CHAT_MODEL})...", flush=True)
        response = await asyncio.to_thread(
            chat_client.chat.completions.create,
            model=CHAT_MODEL,
            messages=messages,
            max_tokens=4096,
            temperature=0.0,
        )
        full_answer = response.choices[0].message.content or ""
        print(f"[RAG Diagnostics] Raw LLM Answer:\n{full_answer}", flush=True)
        
        # 9. Citation validation and cleaning
        cleaned_answer, had_hallucination = validate_and_clean_citations(full_answer, expanded_chunks)
        if had_hallucination:
            print("[RAG Diagnostics] Hallucinated citations detected and cleaned from answer.", flush=True)
            
        print(f"[RAG Diagnostics] Final Cleaned Answer:\n{cleaned_answer}\n", flush=True)

        # Save validated response, parsed sources into cache (Redis / Local Memory) — but
        # never cache a "not found" answer (see agent_stream for why).
        is_not_found = "could not find" in cleaned_answer.lower()
        if not is_not_found:
            _query_cache.set(machine_id, clean_chat_key, clean_msg, cleaned_answer, sources)

        return {
            "status": "done",
            "content": cleaned_answer,
            "sources": sources,
            "searches": searches
        }
        
    except Exception as e:
        err = str(e)
        print(f"[RAG Diagnostics] ERROR in LLM generation: {err}", flush=True)
        if "rate_limit" in err.lower() or "429" in err:
            error_msg = "Rate limit reached. Please wait 1–2 minutes and try again."
        else:
            error_msg = f"API Error: {err}"
            
        return {
            "status": "error",
            "content": error_msg,
            "sources": [],
            "searches": []
        }