import os
import json
from uuid import uuid4
from groq import Groq
from langchain_core.documents import Document

# CHANGED: Using HuggingFaceEndpointEmbeddings for online API inference
from langchain_huggingface import HuggingFaceEndpointEmbeddings
from langchain_chroma import Chroma
from dotenv import load_dotenv
import random
import shutil
from optimized_quiz import OPTIMIZED_QUESTIONS
from chat_resources import *
from datetime import datetime, timedelta

# Load main env
load_dotenv(dotenv_path=".env")

# Load secrets env (this can override duplicates from .env)
load_dotenv(dotenv_path=".secrets.env", override=True)

# Config
GROQ_API_KEY = os.getenv("GROQ_API_KEY")
HF_API_KEY = os.getenv("HF_API_KEY")
DATA_PATH = "data.json"
CHROMA_PATH = "chroma_db"

TEMPERATURE = float(os.getenv("G_TEMPERATURE", 0.7))
MAX_TOKENS = int(os.getenv("G_MAX_TOKENS", 400))
RETRIEVE_K = int(os.getenv("G_RETRIEVE_K", 3))
TOP_P = float(os.getenv("G_TOP_P", 1.0))
MAX_CONVERSATION_HISTORY = int(os.getenv("G_MAX_CONVERSATION_HISTORY", 5))
MMR = str(os.getenv("MMR", "mmr"))
G_FETCH_K = int(os.getenv("G_FETCH_K", 20))
LAMBDA_MULT = float(os.getenv("LAMBDA_MULT", 0.5))


class GroqClient:
    def __init__(self):
        self._sessions = {}  # {ip: {'shown': set(), 'last_activity': datetime}}
        self.SESSION_TIMEOUT = timedelta(minutes=30)

        self.documents = self.load_json_data(DATA_PATH)
        if not self.documents:
            raise RuntimeError("No data loaded")

        self.vector_store = self.init_vector_store(self.documents)

        self.retriever = self.vector_store.as_retriever(
            search_type=MMR,  # Use Maximal Marginal Relevance
            search_kwargs={
                "k": RETRIEVE_K,  # Final number of docs to return
                "fetch_k": G_FETCH_K,  # Number of docs to initially fetch before filtering for diversity
                "lambda_mult": LAMBDA_MULT,  # Balance between relevance (1.0) and diversity (0.0)
            },
        )

        if not GROQ_API_KEY:
            raise RuntimeError("GROQ_API_KEY not found in environment")

        self.client = Groq(api_key=GROQ_API_KEY)

        self.SYSTEM_MESSSAGE = SYSTEM_MESSSAGE

        self.PROMPT_TEMPLATE = PROMPT_TEMPLATE

        self.BLACKLIST = BLACKLIST

    def load_json_data(self, path):
        try:
            with open(path, "r", encoding="utf-8") as f:
                data = json.load(f)

            documents = []

            if "qa" in data:
                for item in data["qa"]:
                    text = f"Q: {item['question']}\nA: {item['answer']}"
                    documents.append(
                        Document(
                            page_content=text,
                            metadata={
                                "id": item.get("id", str(uuid4())),
                                "category": item.get("category", "QA"),
                            },
                        )
                    )

            if "chunks" in data:
                for item in data["chunks"]:
                    documents.append(
                        Document(
                            page_content=item["chunk"],
                            metadata={
                                "id": item.get("id", str(uuid4())),
                                "category": "Chunk",
                            },
                        )
                    )

            return documents

        except Exception as e:
            print(f"Error loading JSON data: {e}")
            return []

    def init_vector_store(self, documents):
        # CHANGED: Using online HuggingFaceEndpointEmbeddings with API key
        # This uses HuggingFace's hosted inference API instead of local model
        if not HF_API_KEY:
            raise RuntimeError("HF_API_KEY not found in environment")

        embeddings_model = HuggingFaceEndpointEmbeddings(
            model="sentence-transformers/all-MiniLM-L6-v2",
            huggingfacehub_api_token=HF_API_KEY,
        )

        # Clear old data to avoid duplicates
        if os.path.exists(CHROMA_PATH):
            shutil.rmtree(CHROMA_PATH)

        uuids = [str(uuid4()) for _ in documents]

        vector_store = Chroma(
            collection_name="user_data",
            embedding_function=embeddings_model,
            persist_directory=CHROMA_PATH,
        )

        # CHANGED: This now makes API calls to HuggingFace for embeddings
        vector_store.add_documents(documents=documents, ids=uuids)
        return vector_store

    def handle_unknown_query(self):
        return random.choice(FALLBACK_RESPONSES)

    # def get_next_questions(self):
    #     return random.sample(OPTIMIZED_QUESTIONS, 3)
    def cleanup_expired_sessions(self):
        """Remove expired sessions to avoid memory overload."""
        current_time = datetime.now()
        expired_ips = [
            ip
            for ip, data in self._sessions.items()
            if current_time - data["last_activity"] > self.SESSION_TIMEOUT
        ]
        for ip in expired_ips:
            del self._sessions[ip]

    def get_next_questions(self, ip: str):
        """Return 3 non-repeated random questions within the session."""
        self.cleanup_expired_sessions()

        current_time = datetime.now()
        if ip not in self._sessions:
            self._sessions[ip] = {"shown": set(), "last_activity": current_time}
        else:
            if (
                current_time - self._sessions[ip]["last_activity"]
                > self.SESSION_TIMEOUT
            ):
                self._sessions[ip]["shown"].clear()
            self._sessions[ip]["last_activity"] = current_time

        shown = self._sessions[ip]["shown"]
        remaining = [q for q in OPTIMIZED_QUESTIONS if q not in shown]

        if len(remaining) < 3:
            shown.clear()
            remaining = OPTIMIZED_QUESTIONS[:]

        selected = random.sample(remaining, 3)
        shown.update(selected)
        return selected

    # ---------------MAIN-----------------

# Non-streaming ask method for backwards compatibility
# ============================================
    def ask(self, raw_query: str) -> tuple[str, str | None]:
        if (
            not raw_query
            or raw_query is None
            or raw_query == ""
            or len(raw_query) > 1000
        ):
            return "Please provide a valid query under 1,000 characters.", None

        if raw_query.lower() in GREETINGS_TRIGGERS:
            return random.choice(GREETINGS), None

        try:
            docs = self.retriever.invoke(raw_query)
        except Exception as e:
            return f"Error retrieving documents: {e}", None

        if not docs:
            return self.handle_unknown_query(), None

        context = "\n".join([d.page_content for d in docs])
        fallback = self.handle_unknown_query()
        prompt = self.PROMPT_TEMPLATE.format(
            context=context, question=raw_query, fallback_response=fallback
        )

        messages = [
            {
                "role": "system",
                "content": self.SYSTEM_MESSSAGE,
            },
        ] + [
            {"role": "user", "content": prompt},
        ]

        # Try multiple models with fallback
        models_to_try = [
            "openai/gpt-oss-120b",
            "openai/gpt-oss-20b",
            "compound-beta-mini",
            "llama-3.1-8b-instant",
            "llama-3.3-70b-versatile",
        ]

        backup_models = [
            "compound-beta-mini",
            "llama-3.1-8b-instant",
            "llama-3.3-70b-versatile",
        ]

        random.shuffle(models_to_try)

        for model in models_to_try:
            try:
                completion = self.client.chat.completions.create(
                    model=model,
                    messages=messages,
                    temperature=TEMPERATURE,
                    max_completion_tokens=MAX_TOKENS,
                    top_p=TOP_P,
                    stream=False,
                )
                response = completion.choices[0].message.content
                if response and response.strip():
                    return response.strip(), model
                else:
                    continue  # Try next model

            except Exception as e:
                # Check if it's a rate limit error
                if "rate_limit_exceeded" in str(e) or "429" in str(e):
                    # print(f"Rate limit hit for model {model}, trying fallback...")
                    continue
                else:
                    # For other errors, return immediately
                    return f"Error while calling LLM: {e}", None

        # If all models fail
        return "I'm temporarily experiencing high demand. Please try again in a few minutes or rephrase your question.", None

# Streaming ask method (Generator)
# ============================================
    def ask_stream(self, raw_query: str):
        """Generator function that yields response chunks in real-time"""

        # Validation
        if not raw_query or raw_query is None or raw_query == "" or len(raw_query) > 1000:
            yield "Please provide a valid query under 1,000 characters."
            return

        if raw_query.lower() in GREETINGS_TRIGGERS:
            yield random.choice(GREETINGS)
            return

        # Retrieve documents
        try:
            docs = self.retriever.invoke(raw_query)
        except Exception as e:
            yield f"Error retrieving documents: {e}"
            return

        if not docs:
            yield self.handle_unknown_query()
            return

        # Prepare context and prompt
        context = "\n".join([d.page_content for d in docs])
        fallback = self.handle_unknown_query()
        prompt = self.PROMPT_TEMPLATE.format(
            context=context, question=raw_query, fallback_response=fallback
        )

        messages = [
            {
                "role": "system",
                "content": self.SYSTEM_MESSSAGE,
            },
            {"role": "user", "content": prompt},
        ]

        # Try multiple models with fallback
        models_to_try = [
            "openai/gpt-oss-120b",
            "openai/gpt-oss-20b",
            "compound-beta-mini",
            "llama-3.1-8b-instant",
            "llama-3.3-70b-versatile",
        ]

        random.shuffle(models_to_try)

        for model in models_to_try:
            try:
                completion = self.client.chat.completions.create(
                    model=model,
                    messages=messages,
                    temperature=TEMPERATURE,
                    max_completion_tokens=MAX_TOKENS,
                    top_p=TOP_P,
                    stream=True,  # Enable streaming
                )

                # Stream chunks as they arrive
                has_content = False
                for chunk in completion:
                    if chunk.choices[0].delta.content:
                        has_content = True
                        yield chunk.choices[0].delta.content

                    # Check if stream finished
                    if chunk.choices[0].finish_reason:
                        break

                # If we got content, we're done
                if has_content:
                    return

            except Exception as e:
                # Check if it's a rate limit error
                if "rate_limit_exceeded" in str(e) or "429" in str(e):
                    continue  # Try next model
                else:
                    yield f"Error while calling LLM: {e}"
                    return

        # If all models fail
        yield "I'm temporarily experiencing high demand. Please try again in a few minutes or rephrase your question."