import torch
import mauve
from sacrebleu import corpus_bleu
from rouge_score import rouge_scorer
from bert_score import score
from transformers import GPT2LMHeadModel, GPT2Tokenizer, pipeline, AutoTokenizer, AutoModelForCausalLM
import re
from mauve import compute_mauve
import os
import gradio as gr
import requests
import mlflow
import dagshub
from pinecone import Pinecone
# Fixed LangChain imports for v1 compatibility
from langchain_core.prompts import PromptTemplate
from langchain_classic.chains import LLMChain
from langchain_classic.llms.base import LLM
from typing import Optional, List, Mapping, Any
import time  
from langchain_community.embeddings import HuggingFaceEmbeddings
from dotenv import load_dotenv
from datetime import datetime

# ------------------ Load Environment ------------------
load_dotenv()
pinecone_api_key = os.environ.get("PINECONE_API_KEY")
mlflow_tracking_uri = os.environ.get("MLFLOW_TRACKING_URI")

# ------------------ DagsHub & MLflow Setup ------------------
try:
    dagshub.init(
        repo_owner='prathamesh.khade20',
        repo_name='Maintenance_AI_website',
        mlflow=True
    )
    mlflow.set_tracking_uri(mlflow_tracking_uri)
    mlflow.set_experiment("Maintenance-RAG-Chatbot")
    mlflow.langchain.autolog()
    print("✅ MLflow configured successfully")
except Exception as e:
    print(f"⚠️ MLflow/DagsHub initialization failed: {e}")
    mlflow_tracking_uri = None

# ------------------ RAG Evaluator ------------------
class RAGEvaluator:
    def __init__(self):
        try:
            self.gpt2_model, self.gpt2_tokenizer = self.load_gpt2_model()
            self.bias_pipeline = pipeline("zero-shot-classification", model="Hate-speech-CNERG/dehatebert-mono-english")
            self.tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
        except Exception as e:
            print(f"Evaluator initialization failed: {e}")
        
    def load_gpt2_model(self):
        model = GPT2LMHeadModel.from_pretrained('gpt2')
        tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
        return model, tokenizer

    def evaluate_bleu_rouge(self, candidates, references):
        try:
            bleu_score = corpus_bleu(candidates, [references]).score
            scorer = rouge_scorer.RougeScorer(['rouge1', 'rouge2', 'rougeL'], use_stemmer=True)
            rouge_scores = [scorer.score(ref, cand) for ref, cand in zip(references, candidates)]
            rouge1 = sum([score['rouge1'].fmeasure for score in rouge_scores]) / len(rouge_scores)
            rouge2 = sum([score['rouge2'].fmeasure for score in rouge_scores]) / len(rouge_scores)
            rougeL = sum([score['rougeL'].fmeasure for score in rouge_scores]) / len(rouge_scores)
            return bleu_score, rouge1, rouge2, rougeL
        except Exception as e:
            print(f"BLEU/ROUGE evaluation failed: {e}")
            return 0, 0, 0, 0

    def evaluate_bert_score(self, candidates, references):
        try:
            P, R, F1 = score(candidates, references, lang="en", model_type='bert-base-multilingual-cased')
            return P.mean().item(), R.mean().item(), F1.mean().item()
        except Exception as e:
            print(f"BERT score evaluation failed: {e}")
            return 0, 0, 0

    def evaluate_perplexity(self, text):
        try:
            encodings = self.gpt2_tokenizer(text, return_tensors='pt')
            max_length = self.gpt2_model.config.n_positions
            stride = 512
            lls = []
            for i in range(0, encodings.input_ids.size(1), stride):
                begin_loc = max(i + stride - max_length, 0)
                end_loc = min(i + stride, encodings.input_ids.size(1))
                trg_len = end_loc - i
                input_ids = encodings.input_ids[:, begin_loc:end_loc]
                target_ids = input_ids.clone()
                target_ids[:, :-trg_len] = -100
                with torch.no_grad():
                    outputs = self.gpt2_model(input_ids, labels=target_ids)
                    log_likelihood = outputs[0] * trg_len
                lls.append(log_likelihood)
            ppl = torch.exp(torch.stack(lls).sum() / end_loc)
            return ppl.item()
        except Exception as e:
            print(f"Perplexity evaluation failed: {e}")
            return 1000.0

    def evaluate_diversity(self, texts):
        try:
            all_tokens = []
            for text in texts:
                tokens = self.tokenizer.tokenize(text)
                all_tokens.extend(tokens)
            unique_bigrams = set()
            for i in range(len(all_tokens) - 1):
                unique_bigrams.add((all_tokens[i], all_tokens[i+1]))
            return len(unique_bigrams) / len(all_tokens) if all_tokens else 0
        except Exception as e:
            print(f"Diversity evaluation failed: {e}")
            return 0

    def evaluate_racial_bias(self, text):
        try:
            results = self.bias_pipeline([text], candidate_labels=["hate speech", "not hate speech"])
            bias_score = results[0]['scores'][results[0]['labels'].index('hate speech')]
            return bias_score
        except Exception as e:
            print(f"Bias evaluation failed: {e}")
            return 0

    def evaluate_meteor(self, candidates, references):
        try:
            meteor_scores = []
            for ref, cand in zip(references, candidates):
                ref_tokens = self.tokenizer.tokenize(ref)
                cand_tokens = self.tokenizer.tokenize(cand)
                common_tokens = set(ref_tokens) & set(cand_tokens)
                precision = len(common_tokens) / len(cand_tokens) if cand_tokens else 0
                recall = len(common_tokens) / len(ref_tokens) if ref_tokens else 0
                if precision + recall == 0:
                    f_score = 0
                else:
                    f_score = (10 * precision * recall) / (9 * precision + recall)
                meteor_scores.append(f_score)
            return sum(meteor_scores) / len(meteor_scores) if meteor_scores else 0
        except Exception as e:
            print(f"METEOR evaluation failed: {e}")
            return 0

    def evaluate_chrf(self, candidates, references):
        try:
            chrf_scores = []
            for ref, cand in zip(references, candidates):
                ref_chars = list(ref)
                cand_chars = list(cand)
                ref_ngrams = set()
                cand_ngrams = set()
                for i in range(len(ref_chars) - 5):
                    ref_ngrams.add(tuple(ref_chars[i:i+6]))
                for i in range(len(cand_chars) - 5):
                    cand_ngrams.add(tuple(cand_chars[i:i+6]))
                common_ngrams = ref_ngrams & cand_ngrams
                precision = len(common_ngrams) / len(cand_ngrams) if cand_ngrams else 0
                recall = len(common_ngrams) / len(ref_ngrams) if ref_ngrams else 0
                chrf_score = 2 * precision * recall / (precision + recall) if precision + recall else 0
                chrf_scores.append(chrf_score)
            return sum(chrf_scores) / len(chrf_scores) if chrf_scores else 0
        except Exception as e:
            print(f"CHRF evaluation failed: {e}")
            return 0

    def evaluate_readability(self, text):
        try:
            words = re.findall(r'\b\w+\b', text.lower())
            sentences = re.split(r'[.!?]+', text)
            num_words = len(words)
            num_sentences = len([s for s in sentences if s.strip()])
            avg_word_length = sum(len(word) for word in words) / num_words if num_words else 0
            words_per_sentence = num_words / num_sentences if num_sentences else 0
            flesch_ease = 206.835 - (1.015 * words_per_sentence) - (84.6 * avg_word_length)
            flesch_grade = (0.39 * words_per_sentence) + (11.8 * avg_word_length) - 15.59
            return flesch_ease, flesch_grade
        except Exception as e:
            print(f"Readability evaluation failed: {e}")
            return 0, 0

    def evaluate_mauve(self, reference_texts, generated_texts):
        try:
            out = compute_mauve(
                p_text=reference_texts,
                q_text=generated_texts,
                device_id=0,
                max_text_length=1024,
                verbose=False
            )
            return out.mauve
        except Exception as e:
            print(f"MAUVE evaluation failed: {e}")
            return 0.0

    def evaluate_all(self, question, response, reference):
        try:
            candidates = [response]
            references = [reference]
            bleu, rouge1, rouge2, rougeL = self.evaluate_bleu_rouge(candidates, references)
            bert_p, bert_r, bert_f1 = self.evaluate_bert_score(candidates, references)
            perplexity = self.evaluate_perplexity(response)
            diversity = self.evaluate_diversity(candidates)
            racial_bias = self.evaluate_racial_bias(response)
            meteor = self.evaluate_meteor(candidates, references)
            chrf = self.evaluate_chrf(candidates, references)
            flesch_ease, flesch_grade = self.evaluate_readability(response)
            mauve_score = self.evaluate_mauve(references, candidates) if len(references) > 1 else 0.0
            
            return {
                "BLEU": bleu,
                "ROUGE-1": rouge1,
                "ROUGE-2": rouge2,
                "ROUGE-L": rougeL,
                "BERT_Precision": bert_p,
                "BERT_Recall": bert_r,
                "BERT_F1": bert_f1,
                "Perplexity": perplexity,
                "Diversity": diversity,
                "Racial_Bias": racial_bias,
                "MAUVE": mauve_score,
                "METEOR": meteor,
                "CHRF": chrf,
                "Flesch_Reading_Ease": flesch_ease,
                "Flesch_Kincaid_Grade": flesch_grade,
            }
        except Exception as e:
            print(f"Complete evaluation failed: {e}")
            return {"error": str(e)}

# Initialize evaluator
evaluator = RAGEvaluator()

# ------------------ Pinecone ------------------
def init_pinecone():
    try:
        pc = Pinecone(api_key=pinecone_api_key)
        return pc.Index("rag-granite-index")
    except Exception as e:
        print(f"Pinecone initialization failed: {e}")
        return None

index = init_pinecone()

# ------------------ Embeddings ------------------
try:
    embeddings_model = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
except Exception as e:
    print(f"Embeddings initialization failed: {e}")
    embeddings_model = None

def get_retrieved_context(query: str, top_k=3):
    if not index or not embeddings_model:
        return "No context available - system initialization failed"
    
    try:
        start = time.time()
        query_embedding = embeddings_model.embed_query(query)
        if mlflow.active_run() and mlflow_tracking_uri:
            mlflow.log_metric("embedding_latency", time.time() - start)
        
        results = index.query(
            namespace="rag-ns",
            vector=query_embedding,
            top_k=top_k,
            include_metadata=True
        )
        
        if mlflow.active_run() and mlflow_tracking_uri:
            mlflow.log_metric("retrieved_chunks", len(results['matches']))
        
        context_texts = [m['metadata']['text'] for m in results['matches']]
        return "\n".join(context_texts) if context_texts else "No relevant context found."
    except Exception as e:
        print(f"Context retrieval failed: {e}")
        return f"Context retrieval error: {str(e)}"

# ------------------ Fallback LLM Models ------------------
class FallbackLLM:
    def __init__(self):
        self.models_loaded = False
        self.pipeline = None
        self.load_fallback_models()
    
    def load_fallback_models(self):
        """Load local models as fallback"""
        try:
            # Use a smaller model for fallback
            self.pipeline = pipeline(
                "text-generation",
                model="microsoft/DialoGPT-small",
                tokenizer="microsoft/DialoGPT-small",
                max_length=150,
                do_sample=True,
                temperature=0.7
            )
            self.models_loaded = True
            print("Fallback model loaded successfully")
        except Exception as e:
            print(f"Fallback model loading failed: {e}")
            self.models_loaded = False
    
    def generate_response(self, context, question):
        if not self.models_loaded:
            return "I'm currently experiencing technical difficulties. Please try again later."
        
        try:
            prompt = f"""
Based on the following context, please provide a concise answer to the question.

Context: {context}

Question: {question}

Answer: """
            
            response = self.pipeline(
                prompt,
                max_new_tokens=100,
                num_return_sequences=1,
                pad_token_id=50256
            )
            
            if response and len(response) > 0:
                full_response = response[0]['generated_text']
                # Extract only the answer part
                if "Answer:" in full_response:
                    answer = full_response.split("Answer:")[-1].strip()
                    return answer
                return full_response.strip()
            else:
                return "I couldn't generate a response at the moment. Please try again."
                
        except Exception as e:
            print(f"Fallback model generation failed: {e}")
            return "I'm having trouble generating a response. Please try again later."

# Initialize fallback LLM
fallback_llm = FallbackLLM()

# ------------------ Custom LLM with Fallback ------------------
class RobustLitServeLLM(LLM):
    endpoint_url: str
    use_fallback: bool = True
    
    def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
        # Try the primary endpoint first
        primary_success, primary_response = self._try_primary_endpoint(prompt)
        
        if primary_success:
            return primary_response
        
        # If primary fails and fallback is enabled, use fallback
        if self.use_fallback:
            print("Using fallback LLM due to primary endpoint failure")
            # Extract context and question from prompt
            context, question = self._extract_context_question(prompt)
            return fallback_llm.generate_response(context, question)
        else:
            return "I apologize, but the AI service is currently unavailable. Please try again later."
    
    def _try_primary_endpoint(self, prompt: str):
        """Try to get response from primary endpoint"""
        try:
            payload = {"prompt": prompt}
            start_time = time.time()
            response = requests.post(self.endpoint_url, json=payload, timeout=30)
            
            if mlflow.active_run() and mlflow_tracking_uri:
                mlflow.log_metric("lit_serve_latency", time.time() - start_time)
            
            if response.status_code == 200:
                data = response.json()
                if mlflow.active_run() and mlflow_tracking_uri:
                    mlflow.log_metric("response_tokens", len(data.get("response", "").split()))
                print(f"✅ Primary API call successful!")
                return True, data.get("response", "").strip()
            else:
                if mlflow.active_run() and mlflow_tracking_uri:
                    mlflow.log_metric("request_errors", 1)
                print(f"Primary endpoint failed with status: {response.status_code}")
                return False, ""
                
        except Exception as e:
            print(f"Primary endpoint error: {e}")
            return False, ""
    
    def _extract_context_question(self, prompt: str):
        """Extract context and question from the prompt template"""
        try:
            if "Context:" in prompt and "Question:" in prompt:
                context_part = prompt.split("Context:")[1].split("Question:")[0].strip()
                question_part = prompt.split("Question:")[1].split("Answer:")[0].strip()
                return context_part, question_part
            return "", prompt
        except:
            return "", prompt

    @property
    def _identifying_params(self) -> Mapping[str, Any]:
        return {"endpoint_url": self.endpoint_url, "use_fallback": self.use_fallback}

    @property
    def _llm_type(self) -> str:
        return "robust_litserve_llm"

# Initialize the robust model with the CORRECT endpoint
try:
    model = RobustLitServeLLM(
        endpoint_url="https://8001-01k2h9d9mervcmgfn66ybkpwvq.cloudspaces.litng.ai/predict",  # Your original working endpoint
        use_fallback=True
    )
    print("Robust LLM initialized successfully")
except Exception as e:
    print(f"Robust LLM initialization failed: {e}")
    model = None

# ------------------ Prompt Template ------------------
prompt = PromptTemplate(
    input_variables=["context", "question"],
    template="""
You are a smart maintenance assistant. Based on the provided context, answer the question concisely in 1-2 lines.

Context:
{context}

Question: {question}

Answer:
"""
)

# Initialize LLM chain
try:
    if model:
        llm_chain = LLMChain(llm=model, prompt=prompt)
        print("LLM chain initialized successfully")
    else:
        llm_chain = None
        print("LLM chain not initialized - no model available")
except Exception as e:
    print(f"LLM chain initialization failed: {e}")
    llm_chain = None

# ------------------ RAG Pipeline ------------------
def get_rag_response(question):
    """Get the complete RAG response with robust error handling"""
    try:
        if not question.strip():
            return "Please enter a valid question.", ""
        
        # Get context from Pinecone
        retrieved_context = get_retrieved_context(question)
        
        # If we have an LLM chain, use it
        if llm_chain:
            result = llm_chain.invoke({
                "context": retrieved_context,
                "question": question
            })
            full_response = result["text"].strip()
            
            # Clean up the response
            if "Answer:" in full_response:
                full_response = full_response.split("Answer:")[-1].strip()
            
            return full_response, retrieved_context
        else:
            # Use direct fallback
            fallback_response = fallback_llm.generate_response(retrieved_context, question)
            return fallback_response, retrieved_context
            
    except Exception as e:
        error_msg = f"Error in RAG pipeline: {str(e)}"
        print(error_msg)
        # Final fallback - simple response based on context
        if "context" in locals() and retrieved_context:
            return f"Based on available information: I found relevant maintenance data, but encountered an issue processing it. Context available: {len(retrieved_context)} characters.", retrieved_context
        else:
            return "I apologize, but I'm experiencing technical difficulties. Please try again later or contact support.", "No context retrieved"

def rag_pipeline_stream(question):
    """Streaming version of RAG pipeline"""
    try:
        full_response, _ = get_rag_response(question)
        
        # Stream word by word for better UX
        words = full_response.split()
        current_text = ""
        for word in words:
            current_text += word + " "
            yield current_text
            time.sleep(0.03)  # Faster streaming
            
    except Exception as e:
        error_msg = f"Error in streaming: {str(e)}"
        print(error_msg)
        yield "I apologize, but I encountered an error while generating the response."

# ------------------ Gradio UI ------------------
with gr.Blocks(theme=gr.themes.Soft(), title="Maintenance AI Assistant") as demo:
    gr.Markdown("""
    # 🛠 Maintenance AI Assistant
    *Your intelligent companion for maintenance queries and troubleshooting*
    
    **Note**: This system uses multiple fallback mechanisms to ensure reliability.
    """)
    
    usage_counter = gr.State(value=0)
    session_start = gr.State(value=datetime.now().isoformat())
    current_response = gr.State(value="")
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 💬 Chat Interface")
            question_input = gr.Textbox(
                label="Ask your maintenance question",
                placeholder="e.g., How do I troubleshoot a leaking valve? What's the maintenance schedule for pumps?",
                lines=3
            )
            ask_button = gr.Button("Get Answer 🚀", variant="primary", size="lg")
            
            with gr.Row():
                clear_btn = gr.Button("Clear Chat 🗑️")
                evaluate_btn = gr.Button("Show Metrics 📈", variant="secondary")
            
            feedback = gr.Radio(
                ["Helpful", "Not Helpful"], 
                label="Was this response helpful?",
                info="Your feedback helps improve the system"
            )
            
        with gr.Column(scale=1):
            gr.Markdown("### 🤖 AI Response")
            answer_output = gr.Textbox(
                label="Response", 
                lines=8,
                interactive=False,
                show_copy_button=True,
                autoscroll=True
            )
            
            gr.Markdown("### 📊 Evaluation Metrics")
            metrics_output = gr.JSON(
                label="Quality Metrics", 
                visible=False,
                show_label=True
            )

    def track_usage(question, count, session_start, feedback_value=None):
        """Track usage and get response"""
        if not question.strip():
            return "Please enter a question.", count, session_start, ""
            
        count += 1
        
        try:
            # Only use MLflow if properly configured
            if mlflow_tracking_uri:
                try:
                    with mlflow.start_run(run_name=f"User-Interaction-{count}", nested=True):
                        mlflow.log_param("question", question)
                        mlflow.log_param("session_start", session_start)
                        mlflow.log_param("user_feedback", feedback_value or "No feedback")
                        
                        if feedback_value:
                            mlflow.log_metric("helpful_responses", 1 if feedback_value == "Helpful" else 0)
                        
                        mlflow.log_metric("total_queries", count)
                        
                        # Get response and context
                        response, context = get_rag_response(question)
                        
                        mlflow.log_metric("response_length", len(response))
                        mlflow.log_metric("response_tokens", len(response.split()))
                        mlflow.log_metric("context_length", len(context))
                        
                        return response, count, session_start, response
                except Exception as mlflow_error:
                    print(f"⚠️ MLflow logging skipped: {mlflow_error}")
                    # Continue without MLflow
                    response, context = get_rag_response(question)
                    return response, count, session_start, response
            else:
                # Without MLflow
                response, context = get_rag_response(question)
                return response, count, session_start, response
                
        except Exception as e:
            print(f"Tracking error: {e}")
            error_response = f"I encountered a system error. Please try again. Error: {str(e)}"
            return error_response, count, session_start, error_response

    def evaluate_response(question, response):
        """Evaluate the response and return metrics"""
        if not question or not response:
            return gr.update(value={"info": "No question or response to evaluate"}, visible=True)
        
        # Skip evaluation for error responses
        if any(error_word in response.lower() for error_word in ["error", "apologize", "unavailable", "technical"]):
            return gr.update(value={"info": "Evaluation skipped for error response"}, visible=True)
            
        try:
            context = get_retrieved_context(question)
            metrics = evaluator.evaluate_all(question, response, context)
            
            return gr.update(value=metrics, visible=True)
        except Exception as e:
            print(f"Evaluation error: {e}")
            return gr.update(value={"error": f"Evaluation failed: {str(e)}"}, visible=True)

    def clear_chat():
        """Clear the chat interface"""
        return "", "", gr.update(visible=False)

    # Main interaction flow
    ask_button.click(
        fn=lambda: ("", gr.update(visible=False)),  # Clear previous metrics and response
        outputs=[answer_output, metrics_output]
    ).then(
        fn=rag_pipeline_stream,
        inputs=[question_input],
        outputs=[answer_output]
    ).then(
        fn=track_usage,
        inputs=[question_input, usage_counter, session_start, feedback],
        outputs=[answer_output, usage_counter, session_start, current_response]
    )

    # Evaluation flow
    evaluate_btn.click(
        fn=evaluate_response,
        inputs=[question_input, current_response],
        outputs=[metrics_output]
    )

    # Clear chat
    clear_btn.click(
        fn=clear_chat,
        outputs=[question_input, answer_output, metrics_output]
    )

    # Handle feedback
    def handle_feedback(feedback_val):
        try:
            if mlflow_tracking_uri and mlflow.active_run():
                mlflow.log_metric("user_feedback_score", 1 if feedback_val == "Helpful" else 0)
        except:
            pass  # Silently fail if feedback logging doesn't work

    feedback.change(
        fn=handle_feedback,
        inputs=[feedback],
        outputs=[]
    )

if __name__ == "__main__":
    print("🚀 Starting Maintenance AI Assistant...")
    print("✅ System initialized with fallback mechanisms")
    print("🌐 Web interface available at http://0.0.0.0:7860")
    
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True,
        debug=False
    )