import os
import getpass
import spacy
import pandas as pd
from typing import Optional, List, Dict, Any
import subprocess
from langchain.llms.base import LLM
from langchain.docstore.document import Document
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS
from langchain.chains import RetrievalQA
from smolagents import CodeAgent, DuckDuckGoSearchTool, ManagedAgent, LiteLLMModel
from pydantic import BaseModel, Field, ValidationError, validator
from mistralai import Mistral
from langchain.prompts import PromptTemplate

# Import chains and tools
from classification_chain import get_classification_chain
from cleaner_chain import get_cleaner_chain
from refusal_chain import get_refusal_chain
from tailor_chain import get_tailor_chain
from prompts import classification_prompt, refusal_prompt, tailor_prompt

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

# Pydantic models for validation and type safety
class QueryInput(BaseModel):
    query: str = Field(..., min_length=1, description="The input query string")
    
    @validator('query')
    def check_query_is_string(cls, v):
        if not isinstance(v, str):
            raise ValueError("Query must be a valid string")
        if v.strip() == "":
            raise ValueError("Query cannot be empty or just whitespace")
        return v.strip()

class ClassificationResult(BaseModel):
    category: str = Field(..., description="The classification category")
    confidence: float = Field(..., ge=0.0, le=1.0, description="Classification confidence score")

class ModerationResult(BaseModel):
    is_safe: bool = Field(..., description="Whether the content is safe")
    categories: Dict[str, bool] = Field(default_factory=dict, description="Detected content categories")
    original_text: str = Field(..., description="The original input text")

class RAGResponse(BaseModel):
    answer: str = Field(..., description="The generated answer")
    sources: List[str] = Field(default_factory=list, description="Source documents used")
    confidence: float = Field(..., ge=0.0, le=1.0, description="Confidence score of the answer")

# Load spaCy model for NER
def install_spacy_model():
    try:
        spacy.load("en_core_web_sm")
        print("spaCy model 'en_core_web_sm' is already installed.")
    except OSError:
        print("Downloading spaCy model 'en_core_web_sm'...")
        subprocess.run(["python", "-m", "spacy", "download", "en_core_web_sm"], check=True)
        print("spaCy model 'en_core_web_sm' downloaded successfully.")

install_spacy_model()
nlp = spacy.load("en_core_web_sm")

def extract_main_topic(query: str) -> str:
    try:
        query_input = QueryInput(query=query)
        doc = nlp(query_input.query)
        main_topic = None
        
        # Try to find named entities first
        for ent in doc.ents:
            if ent.label_ in ["ORG", "PRODUCT", "PERSON", "GPE", "TIME"]:
                main_topic = ent.text
                break
        
        # If no named entities found, look for nouns
        if not main_topic:
            for token in doc:
                if token.pos_ in ["NOUN", "PROPN"]:
                    main_topic = token.text
                    break
        
        return main_topic if main_topic else "this topic"
    except Exception as e:
        print(f"Error extracting main topic: {e}")
        return "this topic"

def moderate_text(query: str) -> ModerationResult:
    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:
            categories = {
                "violence": response.results[0].categories.get("violence_and_threats", False),
                "hate": response.results[0].categories.get("hate_and_discrimination", False),
                "dangerous": response.results[0].categories.get("dangerous_and_criminal_content", False),
                "selfharm": response.results[0].categories.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 e:
        raise ValueError(f"Input validation failed: {str(e)}")
    except Exception as e:
        raise RuntimeError(f"Moderation failed: {str(e)}")

def classify_query(query: str) -> ClassificationResult:
    try:
        query_input = QueryInput(query=query)
        
        wellness_keywords = ["box breathing", "meditation", "yoga", "mindfulness", "breathing exercises"]
        if any(keyword in query_input.query.lower() for keyword in wellness_keywords):
            return ClassificationResult(category="Wellness", confidence=0.9)
        
        class_result = classification_chain.invoke({"query": query_input.query})
        classification = class_result.get("text", "").strip()
        
        confidence_map = {
            "Wellness": 0.8,
            "Brand": 0.8,
            "OutOfScope": 0.6
        }
        
        return ClassificationResult(
            category=classification if classification != "" else "OutOfScope",
            confidence=confidence_map.get(classification, 0.5)
        )
    except ValidationError as e:
        raise ValueError(f"Classification input validation failed: {str(e)}")
    except Exception as e:
        raise RuntimeError(f"Classification failed: {str(e)}")

def build_or_load_vectorstore(csv_path: str, store_dir: str) -> FAISS:
    try:
        if os.path.exists(store_dir):
            print(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)
        
        print(f"Building new FAISS store from CSV: {csv_path}")
        df = pd.read_csv(csv_path)
        df = df.loc[:, ~df.columns.str.contains('^Unnamed')]
        df.columns = df.columns.str.strip()
        
        # Handle column name variations
        if "Answer" in df.columns:
            df.rename(columns={"Answer": "Answers"}, inplace=True)
        if "Question" not in df.columns and "Question " 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 = [
            Document(page_content=str(row["Answers"]), metadata={"question": str(row["Question"])})
            for _, row in df.iterrows()
        ]
        
        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
        
    except Exception as e:
        raise RuntimeError(f"Error building/loading vector store: {str(e)}")

def build_rag_chain(llm_model: LiteLLMModel, vectorstore: FAISS) -> RetrievalQA:
    class GeminiLangChainLLM(LLM):
        def _call(self, prompt: str, stop: Optional[list] = None, **kwargs) -> str:
            messages = [{"role": "user", "content": prompt}]
            return llm_model(messages, stop_sequences=stop)

        @property
        def _llm_type(self) -> str:
            return "custom_gemini"
    
    try:
        retriever = vectorstore.as_retriever(search_type="similarity", search_kwargs={"k": 3})
        gemini_as_llm = GeminiLangChainLLM()
        return RetrievalQA.from_chain_type(
            llm=gemini_as_llm,
            chain_type="stuff",
            retriever=retriever,
            return_source_documents=True
        )
    except Exception as e:
        raise RuntimeError(f"Error building RAG chain: {str(e)}")

def do_web_search(query: str) -> str:
    try:
        query_input = QueryInput(query=query)
        search_tool = DuckDuckGoSearchTool()
        web_agent = CodeAgent(tools=[search_tool], model=pydantic_agent)
        managed_web_agent = ManagedAgent(agent=web_agent, name="web_search", description="Performs web searches")
        manager_agent = CodeAgent(tools=[], model=pydantic_agent, managed_agents=[managed_web_agent])
        
        search_query = f"Give me relevant info: {query_input.query}"
        return manager_agent.run(search_query)
    except Exception as e:
        return f"Web search failed: {str(e)}"

def merge_responses(kb_answer: str, web_answer: str) -> str:
    try:
        if not kb_answer and not web_answer:
            return "No relevant information found."
        
        if not web_answer:
            return kb_answer.strip()
            
        if not kb_answer:
            return web_answer.strip()
            
        return f"Knowledge Base Answer: {kb_answer.strip()}\n\nWeb Search Result: {web_answer.strip()}"
    except Exception as e:
        return f"Error merging responses: {str(e)}"
def sanitize_message(message: Any) -> str:
    """Sanitize message input to ensure it's a valid string."""
    if hasattr(message, 'content'):
        return str(message.content)
    if isinstance(message, (list, dict)):
        return str(message)
    return str(message)

# Modify your run_pipeline function to include the sanitization
def run_pipeline(query: str) -> str:
    try:
        # Sanitize input
        query = sanitize_message(query)
        
        # Rest of your pipeline code...
        moderation_result = moderate_text(query)
        if not moderation_result.is_safe:
            return "Sorry, this query contains harmful or inappropriate content."
        # Validate and moderate input
        moderation_result = moderate_text(query)
        if not moderation_result.is_safe:
            return "Sorry, this query contains harmful or inappropriate content."

        # Classify the query
        classification_result = classify_query(moderation_result.original_text)

        if classification_result.category == "OutOfScope":
            refusal_text = refusal_chain.run({"topic": "this topic"})
            return tailor_chain.run({"response": refusal_text}).strip()

        # Handle different classifications
        if classification_result.category == "Wellness":
            rag_result = wellness_rag_chain({"query": moderation_result.original_text})
            csv_answer = rag_result["result"].strip()
            web_answer = "" if csv_answer else do_web_search(moderation_result.original_text)
            final_merged = merge_responses(csv_answer, web_answer)
            return tailor_chain.run({"response": final_merged}).strip()

        if classification_result.category == "Brand":
            rag_result = brand_rag_chain({"query": moderation_result.original_text})
            csv_answer = rag_result["result"].strip()
            final_merged = merge_responses(csv_answer, "")
            return tailor_chain.run({"response": final_merged}).strip()

        # Default fallback
        refusal_text = refusal_chain.run({"topic": "this topic"})
        return tailor_chain.run({"response": refusal_text}).strip()

    except Exception as e:
        return f"An error occurred while processing your request: {str(e)}"

# Initialize chains and vectorstores
try:
    classification_chain = get_classification_chain()
    refusal_chain = get_refusal_chain()
    tailor_chain = get_tailor_chain()
    cleaner_chain = get_cleaner_chain()

    wellness_csv = "AIChatbot.csv"
    brand_csv = "BrandAI.csv"
    wellness_store_dir = "faiss_wellness_store"
    brand_store_dir = "faiss_brand_store"

    wellness_vectorstore = build_or_load_vectorstore(wellness_csv, wellness_store_dir)
    brand_vectorstore = build_or_load_vectorstore(brand_csv, brand_store_dir)

    gemini_llm = LiteLLMModel(model_id="gemini/gemini-pro", api_key=os.environ.get("GEMINI_API_KEY"))
    wellness_rag_chain = build_rag_chain(gemini_llm, wellness_vectorstore)
    brand_rag_chain = build_rag_chain(gemini_llm, brand_vectorstore)
    
    print("Pipeline initialized successfully!")
except Exception as e:
    print(f"Error initializing pipeline: {str(e)}")

def run_with_chain(query: str) -> str:
    return run_pipeline(query)