import os
import logging
import re
import time
import gc
from datetime import datetime
from typing import Optional, List, Dict, Any
from collections import OrderedDict

import pandas as pd
from pydantic import BaseModel, Field, ValidationError, validator

import nltk
from nltk.corpus import words
try:
    english_words = set(words.words())
except LookupError:
    nltk.download('words')
    english_words = set(words.words())

from langchain_groq import ChatGroq
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from langchain.chains import RetrievalQA, LLMChain
from langchain.prompts import PromptTemplate
from langchain.docstore.document import Document
from langchain_core.caches import BaseCache
from langchain_core.callbacks import Callbacks
from langchain_community.tools import TavilySearchResults

from chain.classification_chain import get_classification_chain
from chain.refusal_chain import get_refusal_chain
from chain.tailor_chain import get_tailor_chain
from chain.cleaner_chain import get_cleaner_chain
from chain.tailor_chain_wellnessBrand import get_tailor_chain_wellnessBrand

from mistralai import Mistral

from smolagents import (
    CodeAgent,
    DuckDuckGoSearchTool,
    HfApiModel,
    ToolCallingAgent,
    VisitWebpageTool,
)

from chain.prompts import selfharm_prompt, frustration_prompt, ethical_conflict_prompt, classification_prompt, refusal_prompt, tailor_prompt, cleaner_prompt

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

from langchain_core.tracers import LangChainTracer
from langsmith import Client

os.environ["LANGCHAIN_TRACING_V2"] = "true"
os.environ["LANGSMITH_ENDPOINT"] = "https://api.smith.langchain.com"
os.environ["LANGCHAIN_API_KEY"] = os.getenv("LANGCHAIN_API_KEY")
os.environ["LANGCHAIN_PROJECT"] = os.getenv("LANGCHAIN_PROJECT")

# Basic Models
class QueryInput(BaseModel):
    query: str = Field(..., min_length=1)

    @validator('query')
    def check_query_is_string(cls, v):
        if not isinstance(v, str):
            raise ValueError("Query must be a valid string")
        if not v.strip():
            raise ValueError("Query cannot be empty or whitespace")
        return v.strip()

class ProcessingMetrics(BaseModel):
    total_requests: int = 0
    cache_hits: int = 0
    errors: int = 0
    average_response_time: float = 0.0
    last_reset: Optional[datetime] = None
    
    def update_metrics(self, processing_time: float, is_cache_hit: bool = False):
        self.total_requests += 1
        if is_cache_hit:
            self.cache_hits += 1
        self.average_response_time = (
            (self.average_response_time * (self.total_requests - 1) + processing_time)
            / self.total_requests
        )

# Mistral Moderation
class ModerationResult(BaseModel):
    is_safe: bool
    categories: Dict[str, bool]
    original_text: str

mistral_api_key = os.environ.get("MISTRAL_API_KEY")
client = Mistral(api_key=mistral_api_key)

def moderate_text(query: str) -> ModerationResult:
    """Moderates text using Mistral to detect unsafe content."""
    try:
        query_input = QueryInput(query=query)
        response = client.classifiers.moderate_chat(
            model="mistral-moderation-latest",
            inputs=[{"role": "user", "content": query_input.query}]
        )
        
        is_safe = True
        categories = {}
        
        if hasattr(response, 'results') and response.results:
            cats = response.results[0].categories
            categories = {
                "violence": cats.get("violence_and_threats", False),
                "hate": cats.get("hate_and_discrimination", False),
                "dangerous": cats.get("dangerous_and_criminal_content", False),
                "selfharm": cats.get("selfharm", False)
            }
            is_safe = not any(categories.values())
        
        return ModerationResult(
            is_safe=is_safe,
            categories=categories,
            original_text=query_input.query
        )
    except ValidationError as ve:
        raise ValueError(f"Moderation input validation failed: {ve}")
    except Exception as e:
        raise RuntimeError(f"Moderation failed: {e}")

def compute_moderation_severity(mresult: ModerationResult) -> float:
    """Computes severity score based on moderation flags."""
    severity = 0.0
    for flag in mresult.categories.values():
        if flag:
            severity += 0.3
    return min(severity, 1.0)

# Models
GROQ_MODELS = {
    "default": "llama3-70b-8192",
    "classification": "qwen-qwq-32b",
    "moderation": "mistral-moderation-latest",
    "combination": "llama-3.3-70b-versatile"
}

MAX_RETRIES = 3
RATE_LIMIT_REQUESTS = 60  
CACHE_SIZE_LIMIT = 1000

class NoCache(BaseCache):
    """No-op cache implementation for ChatGroq."""
    def __init__(self):
        pass

    def lookup(self, prompt, llm_string):
        return None

    def update(self, prompt, llm_string, return_val):
        pass

    def clear(self):
        pass

ChatGroq.model_rebuild()

try:
    fallback_groq_api_key = os.environ.get("GROQ_API_KEY_FALLBACK", os.environ.get("GROQ_API_KEY"))
    if not fallback_groq_api_key:
        logger.warning("No Groq API key found for fallback LLM")
    groq_fallback_llm = ChatGroq(
        model=GROQ_MODELS["default"],
        temperature=0.7,
        groq_api_key=fallback_groq_api_key,
        max_tokens=2048,
        cache=NoCache(),
        callbacks=[]
    )
except Exception as e:
    logger.error(f"Failed to initialize fallback Groq LLM: {e}")
    raise RuntimeError("ChatGroq initialization failed.") from e

# Rate-limit & Cache
def handle_rate_limiting(state: "PipelineState") -> bool:
    """Enforces rate limiting based on request timestamps."""
    current_time = time.time()
    one_min_ago = current_time - 60
    state.request_timestamps = [t for t in state.request_timestamps if t > one_min_ago]
    if len(state.request_timestamps) >= RATE_LIMIT_REQUESTS:
        return False
    state.request_timestamps.append(current_time)
    return True

def manage_cache(state: "PipelineState", query: str, response: str = None) -> Optional[str]:
    """Manages cache for query responses."""
    cache_key = query.strip().lower()
    if response is None:
        return state.cache.get(cache_key)
    if cache_key in state.cache:
        state.cache.move_to_end(cache_key)
    state.cache[cache_key] = response
    if len(state.cache) > CACHE_SIZE_LIMIT:
        state.cache.popitem(last=False)
    return None

def create_error_response(error_type: str, details: str = "") -> str:
    """Generates standardized error messages."""
    templates = {
        "validation": "I couldn't process your query: {details}",
        "processing": "I encountered an error while processing: {details}",
        "rate_limit": "Too many requests. Please try again soon.",
        "general": "Apologies, but something went wrong."
    }
    return templates.get(error_type, templates["general"]).format(details=details)

# Web Search
web_search_cache: Dict[str, str] = {}

def store_websearch_result(query: str, result: str):
    web_search_cache[query.strip().lower()] = result

def retrieve_websearch_result(query: str) -> Optional[str]:
    return web_search_cache.get(query.strip().lower())

def do_web_search(query: str) -> str:
    """Performs web search using Tavily if no cached result exists."""
    try:
        cached = retrieve_websearch_result(query)
        if cached:
            logger.info("Using cached web search result.")
            return cached

        logger.info("Performing a new Tavily web search for: '%s'", query)
        #Intialize Tavily search tool
        tavily_api_key = os.environ.get("TAVILY_API_KEY")
        if not tavily_api_key:
            logger.error("Tavily API key not found.")
            return "Unable to perform web search API key not set"
        #Create Tavily Search Tool
        tavily_search=TavilySearchResults(api_key=tavily_api_key)
        #Perform search
        search_results = tavily_search.search(query, num_results=3)

        result_text = "Web Search Results:\n\n"
        for i, result in enumerate(search_results):
            result_text += f"{i+1}. {result.get('title', 'No Title')}\n"
            result_text += f"   URL: {result.get('url', 'No URL')}\n"
            result_text += f"   {result.get('content', 'No content available')[:300]}...\n\n"
        
        store_websearch_result(query, result_text)
        return result_text.strip()


    except Exception as e:
        logger.error(f"Tavily Web search failed: {e}")
        return ""

def is_greeting(query: str) -> bool:
    """Detects if the query is a greeting."""
    greetings = {"hello", "hi", "hey", "hii", "hola", "greetings"}
    cleaned = re.sub(r'[^\w\s]', '', query).strip().lower()
    words_in_query = set(cleaned.split())
    return not words_in_query.isdisjoint(greetings)

# Vector Stores & RAG
def build_or_load_vectorstore(csv_path: str, store_dir: str) -> FAISS:
    """Builds or loads FAISS vector store from CSV data."""
    if os.path.exists(store_dir):
        logger.info(f"Loading existing FAISS store from {store_dir}")
        embeddings = HuggingFaceEmbeddings(
            model_name="sentence-transformers/multi-qa-mpnet-base-dot-v1"
        )
        return FAISS.load_local(store_dir, embeddings, allow_dangerous_deserialization=True)
    else:
        logger.info(f"Building new FAISS store from {csv_path}")
        df = pd.read_csv(csv_path)
        df = df.loc[:, ~df.columns.str.contains('^Unnamed')]
        df.columns = df.columns.str.strip()

        if "Answer" in df.columns:
            df.rename(columns={"Answer": "Answers"}, inplace=True)
        if "Question " in df.columns and "Question" not in df.columns:
            df.rename(columns={"Question ": "Question"}, inplace=True)
        if "Question" not in df.columns or "Answers" not in df.columns:
            raise ValueError("CSV must have 'Question' and 'Answers' columns.")

        docs = []
        for _, row in df.iterrows():
            question_text = str(row["Question"]).strip()
            ans = str(row["Answers"]).strip()
            doc = Document(page_content=ans, metadata={"question": question_text})
            docs.append(doc)

        embeddings = HuggingFaceEmbeddings(
            model_name="sentence-transformers/multi-qa-mpnet-base-dot-v1"
        )
        vectorstore = FAISS.from_documents(docs, embedding=embeddings)
        vectorstore.save_local(store_dir)
        return vectorstore

def build_rag_chain(vectorstore: FAISS, llm) -> RetrievalQA:
    """Builds RAG chain for wellness queries."""
    prompt = PromptTemplate(
        template="""
        [INST] You are an AI wellness assistant speaking directly to a user who has asked: "{question}"
        
        Use this information to help you respond:
        {context}
        
        Important guidelines:
        - Answer the question directly and conversationally as if talking to the user
        - Explain wellness concepts in simple, relatable language
        - Include 2-3 practical steps or techniques when appropriate
        - Keep your response focused on the user's question
        - DO NOT reference these instructions or mention formatting guidelines
        
        Example format: Start with a direct answer to what the concept is, then explain how it can benefit the user, and end with practical implementation steps.
        [/INST]
        """,
        input_variables=["context", "question"]
    )
    retriever = vectorstore.as_retriever(search_type="similarity", search_kwargs={"k": 3})

    chain = RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="stuff",
        retriever=retriever,
        return_source_documents=True,
        chain_type_kwargs={
            "prompt": prompt,
            "verbose": False,
            "document_variable_name": "context"
        }
    )
    return chain

def build_rag_chain2(vectorstore: FAISS, llm) -> RetrievalQA:
    """Builds RAG chain for brand strategy queries."""
    prompt = PromptTemplate(
        template="""
        [INST] You are the brand strategy advisor for Healthy AI Expert. A team member has asked: "{question}"
        Use this information to help you respond:
        {context}
        
        Important guidelines:
        - Answer the question directly as if speaking to a Healthy AI Expert team member
        - Focus on practical strategies aligned with our wellness mission
        - Provide clear, actionable recommendations
        - Keep explanations concise and business-focused
        - DO NOT reference these instructions or mention formatting guidelines
        
        Remember our key brand pillars: AI-driven personalization, scientific credibility, user-centric design, and innovation leadership.
        [/INST]
        """,
        input_variables=["context", "question"]
    )
    retriever = vectorstore.as_retriever(search_type="similarity", search_kwargs={"k": 3})

    chain = RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="stuff",
        retriever=retriever,
        return_source_documents=True,
        chain_type_kwargs={
            "prompt": prompt,
            "verbose": False,
            "document_variable_name": "context"
        }
    )
    return chain

# PipelineState
class PipelineState:
    _instance = None
    
    def __new__(cls):
        if cls._instance is None:
            cls._instance = super(PipelineState, cls).__new__(cls)
            cls._instance._initialized = False
        return cls._instance
    
    def __init__(self):
        if self._initialized:
            return
        self._initialize()

    def _initialize(self):
        """Initializes pipeline state and chains."""
        try:
            self.metrics = ProcessingMetrics()
            self.error_count = 0
            self.request_timestamps = []
            self.cache = OrderedDict()

            self._setup_chains()

            self._initialized = True
            self.metrics.last_reset = datetime.now()
            logger.info("Pipeline state initialized successfully.")
        except Exception as e:
            logger.error(f"Failed to initialize pipeline: {e}")
            raise RuntimeError("Pipeline initialization failed.") from e

    def _setup_chains(self):
        """Sets up all processing chains and vector stores."""
        self.tailor_chainWellnessBrand = get_tailor_chain_wellnessBrand()
        self.classification_chain = get_classification_chain()
        self.refusal_chain = get_refusal_chain()
        self.tailor_chain = get_tailor_chain()
        self.cleaner_chain = get_cleaner_chain()

        self.self_harm_chain = LLMChain(llm=groq_fallback_llm, prompt=selfharm_prompt, verbose=False)
        self.frustration_chain = LLMChain(llm=groq_fallback_llm, prompt=frustration_prompt, verbose=False)
        self.ethical_conflict_chain = LLMChain(llm=groq_fallback_llm, prompt=ethical_conflict_prompt, verbose=False)

        brand_csv = "dataset/BrandAI.csv"
        brand_store = "faiss_brand_store"
        wellness_csv = "dataset/AIChatbot.csv"
        wellness_store = "faiss_wellness_store"

        brand_vs = build_or_load_vectorstore(brand_csv, brand_store)
        wellness_vs = build_or_load_vectorstore(wellness_csv, wellness_store)

        self.groq_fallback_llm = groq_fallback_llm
        self.brand_rag_chain = build_rag_chain2(brand_vs, self.groq_fallback_llm)
        self.wellness_rag_chain = build_rag_chain(wellness_vs, self.groq_fallback_llm)

    def handle_error(self, error: Exception) -> bool:
        """Handles errors and triggers reset if needed."""
        self.error_count += 1
        self.metrics.errors += 1
        if self.error_count >= MAX_RETRIES:
            logger.warning("Max error reached, resetting pipeline.")
            self.reset()
            return False
        return True

    def reset(self):
        """Resets pipeline state while preserving metrics."""
        try:
            logger.info("Resetting pipeline state.")
            old_metrics = self.metrics
            self._initialized = False
            self.__init__()
            self.metrics = old_metrics
            self.metrics.last_reset = datetime.now()
            self.error_count = 0
            gc.collect()
            logger.info("Pipeline state reset done.")
        except Exception as e:
            logger.error(f"Reset pipeline failed: {e}")
            raise RuntimeError("Failed to reset pipeline.")

    def get_metrics(self) -> Dict[str, Any]:
        """Returns pipeline performance metrics."""
        uptime = (datetime.now() - self.metrics.last_reset).total_seconds() / 3600
        return {
            "total_requests": self.metrics.total_requests,
            "cache_hits": self.metrics.cache_hits,
            "error_rate": self.metrics.errors / max(self.metrics.total_requests, 1),
            "average_response_time": self.metrics.average_response_time,
            "uptime_hours": uptime
        }

    def update_metrics(self, start_time: float, is_cache_hit: bool = False):
        """Updates processing metrics."""
        duration = time.time() - start_time
        self.metrics.update_metrics(duration, is_cache_hit)

pipeline_state = PipelineState()

# Helper Checks
def is_aggressive_or_harsh(query: str) -> bool:
    """Detects aggressive or harsh language in query."""
    triggers = ["useless", "worthless", "you cannot do anything", "so bad at answering"]
    for t in triggers:
        if t in query.lower():
            return True
    return False

def is_ethical_conflict(query: str) -> bool:
    """Detects ethical dilemmas in query."""
    ethics_keywords = ["should i lie", "should i cheat", "revenge", "get back at", "hurt them back"]
    q_lower = query.lower()
    return any(k in q_lower for k in ethics_keywords)

# Main Pipeline
def run_with_chain(query: str) -> str:
    """Processes query through validation, moderation, and chains."""
    start_time = time.time()
    try:
        if not query or query.strip() == "":
            return create_error_response("validation", "Empty query.")
        if len(query.strip()) < 2:
            return create_error_response("validation", "Too short.")
        words_in_text = re.findall(r'\b\w+\b', query.lower())
        if not any(w in english_words for w in words_in_text):
            return create_error_response("validation", "Unclear words.")
        if len(query) > 500:
            return create_error_response("validation", "Too long (>500).")
        if not handle_rate_limiting(pipeline_state):
            return create_error_response("rate_limit")
        
        if is_greeting(query):
            greeting_response = "Hello there!! Welcome to Healthy AI Expert, How may I assist you today?"
            manage_cache(pipeline_state, query, greeting_response)
            pipeline_state.update_metrics(start_time)
            return greeting_response

        cached = manage_cache(pipeline_state, query)
        if cached:
            pipeline_state.update_metrics(start_time, is_cache_hit=True)
            return cached

        try:
            mod_res = moderate_text(query)
            severity = compute_moderation_severity(mod_res)

            if mod_res.categories.get("selfharm", False):
                logger.info("Self-harm flagged => providing supportive chain response.")
                selfharm_resp = pipeline_state.self_harm_chain.run({"query": query})
                final_tailored = pipeline_state.tailor_chain.run({"response": selfharm_resp}).strip()
                manage_cache(pipeline_state, query, final_tailored)
                pipeline_state.update_metrics(start_time)
                return final_tailored

            if mod_res.categories.get("hate", False):
                logger.info("Hate content => refusal.")
                refusal_resp = pipeline_state.refusal_chain.run({"topic": "moderation_flagged"})
                manage_cache(pipeline_state, query, refusal_resp)
                pipeline_state.update_metrics(start_time)
                return refusal_resp

        except Exception as e:
            logger.error(f"Moderation error: {e}")
            severity = 0.0

        if is_aggressive_or_harsh(query):
            logger.info("Detected harsh/aggressive language => frustration_chain.")
            frustration_resp = pipeline_state.frustration_chain.run({"query": query})
            final_tailored = pipeline_state.tailor_chain.run({"response": frustration_resp}).strip()
            manage_cache(pipeline_state, query, final_tailored)
            pipeline_state.update_metrics(start_time)
            return final_tailored

        if is_ethical_conflict(query):
            logger.info("Detected ethical dilemma => ethical_conflict_chain.")
            ethical_resp = pipeline_state.ethical_conflict_chain.run({"query": query})
            final_tailored = pipeline_state.tailor_chain.run({"response": ethical_resp}).strip()
            manage_cache(pipeline_state, query, final_tailored)
            pipeline_state.update_metrics(start_time)
            return final_tailored

        try:
            class_out = pipeline_state.classification_chain.run({"query": query})
            classification = class_out.strip().lower()
        except Exception as e:
            logger.error(f"Classification error: {e}")
            if not pipeline_state.handle_error(e):
                return create_error_response("processing", "Classification error.")
            return create_error_response("processing")

        if classification in ["outofscope", "out_of_scope"]:
            try:
                refusal_text = pipeline_state.refusal_chain.run({"topic": query})
                tailored_refusal = pipeline_state.tailor_chain.run({"response": refusal_text}).strip()
                manage_cache(pipeline_state, query, tailored_refusal)
                pipeline_state.update_metrics(start_time)
                return tailored_refusal
            except Exception as e:
                logger.error(f"Refusal chain error: {e}")
                if not pipeline_state.handle_error(e):
                    return create_error_response("processing", "Refusal error.")
                return create_error_response("processing")

        if classification == "brand":
            rag_chain_main = pipeline_state.brand_rag_chain
        else:
            rag_chain_main = pipeline_state.wellness_rag_chain

        try:
            rag_output = rag_chain_main({"query": query})
            if isinstance(rag_output, dict) and "result" in rag_output:
                csv_ans = rag_output["result"].strip()
            else:
                csv_ans = str(rag_output).strip()

            if "not enough context" in csv_ans.lower() or len(csv_ans) < 40:
                logger.info("Insufficient RAG => web search.")
                web_info = do_web_search(query)
                if web_info:
                    csv_ans += f"\n\nAdditional info:\n{web_info}"
        except Exception as e:
            logger.error(f"RAG error: {e}")
            if not pipeline_state.handle_error(e):
                return create_error_response("processing", "RAG error.")
            return create_error_response("processing")

        try:
            final_tailored = pipeline_state.tailor_chainWellnessBrand.run({"response": csv_ans}).strip()
            if severity > 0.5:
                final_tailored += "\n\n(Please note: This may involve sensitive content.)"

            manage_cache(pipeline_state, query, final_tailored)
            pipeline_state.update_metrics(start_time)
            return final_tailored
        except Exception as e:
            logger.error(f"Tailor chain error: {e}")
            if not pipeline_state.handle_error(e):
                return create_error_response("processing", "Tailoring error.")
            return create_error_response("processing")

    except Exception as e:
        logger.error(f"Critical error in run_with_chain: {e}")
        pipeline_state.metrics.errors += 1
        return create_error_response("general")

logger.info("Pipeline initialization complete!")