import logging
import os
import re
from functools import lru_cache
from pathlib import Path
from typing import TypedDict

from dotenv import load_dotenv
from groq import Groq
from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_core.documents import Document
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langgraph.graph import END, START, StateGraph
load_dotenv()
logger = logging.getLogger(__name__)

BASE_DIR = Path(__file__).resolve().parent
FAQ_PATH = BASE_DIR / "support_faq.txt"

EMBEDDING_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
GROQ_MODEL_NAME = os.getenv("GROQ_MODEL", "openai/gpt-oss-20b")

TOP_K = 3
MAX_NEW_TOKENS = 120
MAX_RETRIES = 1

# Similarity score threshold for FAISS distance strategy:
# lower L2 distance means a closer match. This value is a starting point
# and should be checked against the user's own FAQ test questions.
MAX_FAISS_DISTANCE = 1.25

FALLBACK_ANSWER = (
    "I couldn't find a reliable answer in the available FAQ. "
    "Please contact the support team through the Help section. "
    "Never share your password, OTP, or full payment-card details."
)


class SupportState(TypedDict, total=False):
    question: str
    intent: str
    search_query: str
    context: str
    sources: list[str]
    draft: str
    retry_count: int
    final_answer: str
    relevant: bool


class SupportRAG:
    """Initialize and manage the complete RAG pipeline."""

    def __init__(self):
        if not FAQ_PATH.exists():
            raise FileNotFoundError(
                f"Cannot find {FAQ_PATH}. Place support_faq.txt beside support_rag.py."
            )

        faq_text = FAQ_PATH.read_text(encoding="utf-8").strip()
        if not faq_text:
            raise ValueError("support_faq.txt is empty.")

        logger.info("Loading FAQ and building FAISS index.")

        document = Document(
            page_content=faq_text,
            metadata={"source": FAQ_PATH.name},
        )
        splitter = RecursiveCharacterTextSplitter(
            chunk_size=450,
            chunk_overlap=60,
        )
        chunks = splitter.split_documents([document])
        if not chunks:
            raise ValueError("No text chunks were created from support_faq.txt.")

        # Force embeddings onto CPU. In ZeroGPU Spaces, Sentence Transformers
        # may incorrectly auto-detect the emulated CUDA device and fail during
        # CUDA initialization. Groq handles text generation remotely, so CPU
        # embeddings are sufficient for this small FAQ index.
        self.embedding_model = HuggingFaceEmbeddings(
            model_name=EMBEDDING_MODEL_NAME,
            model_kwargs={"device": "cpu"},
        )
        self.vector_store = FAISS.from_documents(chunks, self.embedding_model)
        self.retriever = self.vector_store.as_retriever(
            search_type="similarity",
            search_kwargs={"k": TOP_K},
        )

        # The language model runs through Groq's hosted API.
        # Configure GROQ_API_KEY as a Hugging Face Space Secret.
        api_key = os.getenv("GROQ_API_KEY")
        if not api_key:
            raise RuntimeError(
                "GROQ_API_KEY is missing. Add it under Space Settings > Secrets."
            )
        self.groq_client = Groq(api_key=api_key)
        logger.info("Using Groq model: %s", GROQ_MODEL_NAME)

        self.graph = self._build_graph()
        logger.info("Customer Support RAG pipeline initialized.")

    @staticmethod
    def _normalise(text: str) -> str:
        """Lowercase text and collapse punctuation/whitespace."""
        return re.sub(r"\s+", " ", re.sub(r"[^a-z0-9\s]", " ", text.lower())).strip()

    @staticmethod
    def _classify_question(question: str) -> str:
        """Classify intent using simple keyword rules."""
        text = question.lower()

        categories = {
            "order_tracking": [
                "track", "tracking", "shipment", "delivery",
                "where is my order", "order status", "dispatch",
            ],
            "returns_refunds": [
                "return", "refund", "money back", "exchange",
            ],
            "account_access": [
                "password", "sign in", "sign-in", "login",
                "log in", "account", "otp", "verification code",
            ],
            "billing_payments": [
                "payment", "charged", "charge", "billing",
                "invoice", "card", "transaction", "double charge",
            ],
            "technical_support": [
                "website", "site", "error", "loading",
                "crash", "technical", "not working", "browser", "cache",
            ],
        }

        for category, keywords in categories.items():
            if any(keyword in text for keyword in keywords):
                return category
        return "general_support"

    def classify_intent(self, state: SupportState) -> SupportState:
        question = state["question"].strip()
        return {
            **state,
            "question": question,
            "intent": self._classify_question(question),
            "search_query": question,
            "retry_count": 0,
            "sources": [],
            "relevant": False,
        }

    def retrieve_context(self, state: SupportState) -> SupportState:
        """Retrieve the best FAQ chunks and assess their distance."""
        query = state.get("search_query") or state["question"]

        # FAISS returns L2 distance for this vector store; smaller is closer.
        results = self.vector_store.similarity_search_with_score(query, k=TOP_K)

        context_parts = []
        sources = []
        best_distance = float("inf")

        for doc, distance in results:
            text = doc.page_content.strip()
            best_distance = min(best_distance, float(distance))
            if text:
                context_parts.append(text)
            source = doc.metadata.get("source", "unknown")
            if source not in sources:
                sources.append(source)

        context = "\n\n".join(context_parts)
        relevant = bool(context) and best_distance <= MAX_FAISS_DISTANCE

        logger.info(
            "Retrieved %d chunks for query %r; best FAISS distance=%.4f; relevant=%s",
            len(results), query, best_distance, relevant,
        )

        return {
            **state,
            "context": context,
            "sources": sources if relevant else [],
            "relevant": relevant,
        }

    def _generate(self, prompt: str) -> str:
        """Generate a concise answer using Groq's hosted chat-completions API."""
        messages = [
            {
                "role": "system",
                "content": (
                    "You are a customer support assistant. Answer only from the FAQ "
                    "context provided. Do not use outside knowledge, invent policies, "
                    "or claim access to live systems. If the context does not answer "
                    "the question, say you cannot confirm it. Never request passwords, "
                    "OTPs, or full card details. Give a concise answer in 2 to 4 sentences."
                ),
            },
            {"role": "user", "content": prompt},
        ]

        response = self.groq_client.chat.completions.create(
            model=GROQ_MODEL_NAME,
            messages=messages,
            temperature=0,
            max_tokens=MAX_NEW_TOKENS,
        )
        content = response.choices[0].message.content
        return (content or "").strip()

    def draft_answer(self, state: SupportState) -> SupportState:
        """Generate an answer only if retrieval passed the relevance check."""
        context = state.get("context", "").strip()
        if not context or not state.get("relevant", False):
            return {**state, "draft": "", "final_answer": ""}

        prompt = (
            f"Customer question:\n{state['question']}\n\n"
            f"FAQ context:\n{context}\n\n"
            "Instructions:\n"
            "- Answer the question directly using only the FAQ context.\n"
            "- For returns, explain the Return Request steps and relevant eligibility.\n"
            "- If provided, include the refund processing time.\n"
            "- Do not claim to know live order, refund, or payment status.\n"
            "- If the FAQ does not specify the answer, say it is not specified.\n"
            "- Keep the answer concise, in 2 to 4 short sentences.\n"
            "Answer:"
        )

        try:
            answer = self._generate(prompt)
        except Exception:
            logger.exception("Answer generation failed.")
            answer = ""

        answer = answer.strip()
        logger.info("Generated draft length=%d", len(answer))
        return {**state, "draft": answer}

    def review_answer(self, state: SupportState) -> SupportState:
        """Accept a draft only when retrieval passed the relevance check."""
        answer = state.get("draft", "").strip()
        context = state.get("context", "").strip()

        if answer and context and state.get("relevant", False):
            return {**state, "final_answer": answer}

        return {**state, "final_answer": ""}

    def route_after_review(self, state: SupportState) -> str:
        # An irrelevant query should go directly to fallback rather than retrying
        # with a broader query that could retrieve unrelated chunks.
        if not state.get("relevant", False):
            return "fallback"
        if state.get("final_answer"):
            return "finish"
        if state.get("retry_count", 0) < MAX_RETRIES:
            return "retry"
        return "fallback"

    def rewrite_query(self, state: SupportState) -> SupportState:
        """Add intent-specific terms, then retrieve again."""
        additions = {
            "order_tracking": "order tracking shipment delivery status dispatch",
            "returns_refunds": "refund request return policy eligibility refund processing time",
            "account_access": "account login password recovery sign in reset email",
            "billing_payments": "billing payment duplicate charge transaction",
            "technical_support": "website technical problem troubleshooting browser",
            "general_support": "customer support FAQ policy instructions",
        }
        question = state["question"]
        extra_terms = additions.get(
            state.get("intent", "general_support"),
            additions["general_support"],
        )
        new_state = {
            **state,
            "search_query": f"{question} {extra_terms}",
            "retry_count": state.get("retry_count", 0) + 1,
            "final_answer": "",
            "draft": "",
        }
        return self.retrieve_context(new_state)

    def fallback_answer(self, state: SupportState) -> SupportState:
        return {**state, "final_answer": FALLBACK_ANSWER, "sources": []}

    def _build_graph(self):
        workflow = StateGraph(SupportState)
        workflow.add_node("classify", self.classify_intent)
        workflow.add_node("retrieve", self.retrieve_context)
        workflow.add_node("draft", self.draft_answer)
        workflow.add_node("review", self.review_answer)
        workflow.add_node("retry", self.rewrite_query)
        workflow.add_node("fallback", self.fallback_answer)

        workflow.add_edge(START, "classify")
        workflow.add_edge("classify", "retrieve")
        workflow.add_conditional_edges(
            "retrieve",
            lambda state: "draft" if state.get("relevant", False) else "fallback",
            {"draft": "draft", "fallback": "fallback"},
        )
        workflow.add_edge("draft", "review")
        workflow.add_conditional_edges(
            "review",
            self.route_after_review,
            {"finish": END, "retry": "retry", "fallback": "fallback"},
        )
        workflow.add_edge("retry", "draft")
        workflow.add_edge("fallback", END)
        return workflow.compile()

    def ask(self, question: str) -> dict:
        result = self.graph.invoke({"question": question})
        return {
            "intent": result.get("intent", "general_support"),
            "answer": result.get("final_answer") or FALLBACK_ANSWER,
            "sources": result.get("sources", []),
            "retry_count": result.get("retry_count", 0),
        }


@lru_cache(maxsize=1)
def get_support_rag() -> SupportRAG:
    """Initialize the pipeline once and reuse it."""
    return SupportRAG()


def ask_support(question: str) -> dict:
    """Entry point used by app.py."""
    return get_support_rag().ask(question)