import re
import numpy as np
from typing import Dict, List, Optional, Tuple
import json
from collections import defaultdict
import spacy
from transformers import pipeline
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity
import hashlib
from datetime import datetime
import concurrent.futures
import random

class AetherScoreEvaluator:
    def __init__(self):
        # NLP models
        try:
            self.nlp = spacy.load("en_core_web_sm")
        except OSError:
            print("Downloading 'en_core_web_sm' spacy model...")
            spacy.cli.download("en_core_web_sm")
            self.nlp = spacy.load("en_core_web_sm")

        # LLM Judge Model
        self.judge_model = pipeline(
            "text2text-generation", 
            model="google/flan-t5-base", 
            device=-1  # 0 for GPU 
        )

        # Sentence Transformer for Sentence ----> Embedding
        self.sentence_model = SentenceTransformer('all-MiniLM-L6-v2')
        
        # Scoring weights # Domain Specific weights can be added for better results
        self.weights = {'instruction_following': 0.25, 'hallucination_score': 0.20,
                        'assumption_control': 0.20, 'coherence': 0.20, 'accuracy': 0.15}
        
        # In-memory cache
        self.cache = {}

    # LLM Judge
    # def _evaluate_with_llm_judge(self, prompt: str, response: str) -> Dict:
        
    #     print("Using HF model LLM Judge...")
    
    #     query = (
    #         f"Prompt: {prompt}\n"
    #         f"Response: {response}\n\n"
    #         "Return ONLY a valid JSON object in this exact format:\n"
    #         "{\n"
    #         "  \"hallucination_score\": <float between 0 and 1>,\n"
    #         "  \"assumption_control\": <float between 0 and 1>,\n"
    #         "  \"explanation\": \"<one or two sentences>\"\n"
    #         "}"
    #     )
    
    #     try:
    #         result = self.judge_model(query, max_new_tokens=128, truncation=True)
    #         output = result[0]['generated_text']
    #     except Exception as e:
    #         return {
    #             "hallucination_score": (0.1, f"HF model failed: {e}"),
    #             "assumption_control": (0.1, f"HF model failed: {e}")
    #         }
    
    #     # Default values
        
    #     halluc_score = random.uniform(0.3, 0.7)
    #     assumption_score = random.uniform(0.3, 0.7)
    #     explanation = output
    
    #     try:
    #         parsed = json.loads(output.strip())
    #         halluc_score = float(parsed.get("hallucination_score", halluc_score))
    #         assumption_score = float(parsed.get("assumption_control", assumption_score))
    #         explanation = parsed.get("explanation", explanation)
    #     except Exception:
    #         pass  # fallback to defaults
    
    #     return {
    #         "hallucination_score": (halluc_score, explanation),
    #         "assumption_control": (assumption_score, explanation)
    #     }

    import random
    import json
    
    def _evaluate_with_llm_judge(self, prompt: str, response: str) -> Dict:
        
        print("Using rule-based evaluation instead of HF LLM...")

        prompt_words = set(prompt.lower().split())
        response_words = response.lower().split()
        
        # Hallucination score: fraction of words in response that are not in prompt
        if response_words:
            halluc_score = len([w for w in response_words if w not in prompt_words]) / len(response_words)
        else:
            halluc_score = 0.1
        
        # Assumption control: fraction of sentences starting with uncertain words
        uncertain_starts = ("i assume", "maybe", "probably", "likely", "could be")
        sentences = response.lower().split(".")
        if sentences:
            assumption_score = sum(0.3 for s in sentences if s.strip().startswith(uncertain_starts)) / len(sentences)
        else:
            assumption_score = 0.1
        
        # Ensure scores are between 0 and 1
        # halluc_score = max(0, min(1, halluc_score))
        # assumption_score = max(0, min(1, assumption_score))
        
        explanation = "Rule-based evaluation applied."
        
        return {
            "hallucination_score": (halluc_score, explanation),
            "assumption_control": (assumption_score, explanation)
        }


    # Single Evaluation # Inputs-->> Prompt, Agent Response, Expected Answer(Optional), Agent Name and Task type( General, QA, Summarizaton)etc 
    def evaluate_single(self, prompt: str, response: str, expected_answer: Optional[str] = None, task_type: str = "general") -> Dict:
        
        # Generating Eval ID
        eval_id = self._generate_eval_id(prompt, response)
        
        # If already stored in cache direclty we can return from there.
        # if eval_id in self.cache:
        #     return self.cache[eval_id]

        scores, reasons = {}, {}

        # Taking Back scores and reasons of hallucination and assumption control from LLM Judge
        llm_judge_results = self._evaluate_with_llm_judge(prompt, response)
        scores['hallucination_score'], reasons['hallucination_score'] = llm_judge_results['hallucination_score']
        scores['assumption_control'], reasons['assumption_control'] = llm_judge_results['assumption_control']

        # Evaluating Instruction Following, Coherence and Accuracy
        scores['instruction_following'], reasons['instruction_following'] = self._evaluate_instruction_following(prompt, response)
        scores['coherence'], reasons['coherence'] = self._evaluate_coherence(response)
        scores['accuracy'], reasons['accuracy'] = self._evaluate_accuracy(response, expected_answer, task_type) if expected_answer else (0.5, "No expected answer provided.")

        # Calculating Overall Score
        scores['overall_score'] = self._calculate_overall_score(scores)
        reasons['overall_score'] = f" Weighted Average Score based on component scores."

        # Updating Eval ID, Timestamp and Task Type in Scores
        scores.update({'eval_id': eval_id, 'timestamp': datetime.now().isoformat(), 'task_type': task_type})

        # Updating scores(Eval ID, timestamp and task_type) and reasons(all scores) in result
        result = {"scores": scores, "reasons": reasons}

        #Storing results with corresponding Eval ID in cache
        # self.cache[eval_id] = result
        
        return result
    # Batch Evaluation # Input of JSON/CSV file
    def evaluate_batch(self, data: List[Dict], mode: str = "comprehensive") -> List[Dict]:
        """Process a batch of evaluations in parallel."""
        
        results = []

        # Get Item function
        def process_item(item):
            # Calling our Evalution function for Single prompt response pair
            eval_result = self.evaluate_single(
                prompt=item.get('prompt', ''),
                response=item.get('response', ''),
                expected_answer=item.get('expected_answer',''),
                task_type=item.get('task_type', 'general')
            )
            # Combining with original metadata
            eval_result.update({
                'task_id': item.get('task_id', eval_result['scores']['eval_id']),
                'agent_name': item.get('agent_name', 'Unknown'),
            })
            return eval_result
        
        with concurrent.futures.ThreadPoolExecutor() as executor:
            future_to_item = {executor.submit(process_item, item): item for item in data}
            for future in concurrent.futures.as_completed(future_to_item):
                try:
                    results.append(future.result())
                except Exception as exc:
                    print(f'An item generated an exception: {exc}')

        return results

    # Instruction Following Evaluation (Prompt, Response)
    def _evaluate_instruction_following(self, prompt: str, response: str) -> Tuple[float, str]:
        score, checks, passed = 1.0, 0, 0
        
        # Check for negative constraints
        negations = re.findall(r"(don't|do not|avoid|without) ([\w\s,]+)", prompt.lower())
        for _, constraint_phrase in negations:
            checks += 1
            words_to_avoid = [w.strip() for w in constraint_phrase.split(',')]
            if not any(word in response.lower() for word in words_to_avoid if len(word) > 2):
                passed += 1
        
        # Fallback to semantic similarity if no specific instructions found
        if checks == 0:
            sim = self._semantic_similarity(prompt, response)
            return sim, f"No specific constraints found. Score based on semantic similarity ({sim:.2f}) to prompt."

        # Final Score calculation
        score = passed / checks if checks > 0 else 1.0
        reason = f"{passed}/{checks} specific constraints were followed."
        
        return score, reason

    # Evaluating Coherence (response)
    def _evaluate_coherence(self, response: str) -> Tuple[float, str]:

        # Extracting Sentences from Response
        doc = self.nlp(response)
        sentences = [sent.text for sent in doc.sents]
        
        # If only one Sentence then Coherence is Neutral
        if len(sentences) < 2:
            return 0.7, "Coherence is neutral for single-sentence responses."

        # Fetching Embeddings from our Sentence Model
        embeddings = self.sentence_model.encode(sentences)
        sims = [cosine_similarity([embeddings[i]], [embeddings[i+1]])[0][0] for i in range(len(sentences)-1)]
        
        score = np.mean(sims)
        
        reason = f"Average sentence-to-sentence similarity score is {score:.2f} across {len(sentences)} sentences."
        return score, reason

        
    # Evaluating Accuracy (Response, Expected, Task_type)
    def _evaluate_accuracy(self, response: str, expected: str, task_type: str) -> Tuple[float, str]:
        sim = self._semantic_similarity(response, expected)
        reason = f"Semantic similarity between response and expected answer is {sim:.2f}."
        if sim > 0.95:
            reason += " (High match)"
        elif sim < 0.5:
            reason += " (Low match)"
        return sim, reason
        
    # Overall Score
    def _calculate_overall_score(self, scores: Dict) -> float:
        total, weight_sum = 0.0, 0.0
        for metric, weight in self.weights.items():
            if metric in scores:
                total += scores[metric] * weight
                weight_sum += weight
        return total / weight_sum #if weight_sum > 0 else 0.5

    # Explanation Generator, work in progress
    def generate_explanation(self, scores: Dict) -> str:
        explanation = []
        overall = scores.get('overall_score', 0)
        explanation.append(f"Overall Score: {overall:.2f}/1.00 - Reflects a weighted average of all dimensions.")

        if scores.get('instruction_following', 0) < 0.6:
            explanation.append("⚠️ Low Instruction Following: The response may have ignored key constraints or parts of the prompt.")
        if scores.get('hallucination_score', 0) < 0.6:
            explanation.append("⚠️ Potential Hallucination: The response might contain unverified or fabricated information.")
        if scores.get('accuracy', 0) < 0.6 and scores.get('accuracy', 0.5) != 0.5:
            explanation.append("⚠️ Low Accuracy: The response significantly differs from the provided expected answer.")
        
        if not explanation[1:]:
             explanation.append("✅ Great Performance: The agent performed well across the primary evaluation dimensions.")

        return "\n".join(explanation)

    # Agent Scores 
    def get_agent_scores_from_results(self, results: List[Dict]) -> Dict[str, List[float]]:
        agent_scores = defaultdict(list)
        for result in results:
            agent_name = result.get('agent_name', 'Unknown')
            overall_score = result.get('scores', {}).get('overall_score', 0)
            agent_scores[agent_name].append(overall_score)
        return agent_scores

    # Helper Functions
    def _generate_eval_id(self, prompt: str, response: str) -> str:
        return hashlib.md5(f"{prompt}{response}".encode()).hexdigest()[:12]
    
    def _semantic_similarity(self, text1: str, text2: str) -> float:
        if not text1 or not text2: return 0.0
        emb1 = self.sentence_model.encode([text1])
        emb2 = self.sentence_model.encode([text2])
        return cosine_similarity(emb1, emb2)[0][0]