import os
import requests
import random
import time
import logging
from dotenv import load_dotenv
from messages import krishna_blessings, ayush_teasing, keyword_groups, get_contextual_response, generate_follow_up, handle_vague_input
from ayush_messages import ayush_surprises
from sentence_transformers import SentenceTransformer, util
import joblib
import numpy as np

# Configure logging
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)

# Load environment variables
load_dotenv()
HUGGINGFACE_API_TOKEN = os.getenv("HUGGINGFACE_API_TOKEN")
if not HUGGINGFACE_API_TOKEN:
    logger.error("HUGGINGFACE_API_TOKEN not found in environment variables")
    raise ValueError("HUGGINGFACE_API_TOKEN is required")

# Lazy load sentence transformer model and embeddings
semantic_model = None
keyword_embeddings_cache = None

def init_semantic_model():
    global semantic_model, keyword_embeddings_cache
    if semantic_model is None:
        try:
            semantic_model = SentenceTransformer('all-MiniLM-L6-v2')
            keyword_embeddings_cache = joblib.load('embeddings_cache.joblib')
            logger.debug("Successfully loaded semantic model and embeddings cache")
        except Exception as e:
            logger.error(f"Failed to load semantic model or embeddings: {str(e)}")
            # Retry once
            try:
                semantic_model = SentenceTransformer('all-MiniLM-L6-v2')
                keyword_embeddings_cache = joblib.load('embeddings_cache.joblib')
                logger.debug("Retry successful")
            except Exception as e:
                logger.error(f"Retry failed: {str(e)}")
                semantic_model = None
                keyword_embeddings_cache = {}

# AI model for fallback responses
AI_MODELS = [
    {
        "name": "mistralai/Mixtral-8x7B-Instruct-v0.1",
        "endpoint": "https://api-inference.huggingface.co/models/mistralai/Mixtral-8x7B-Instruct-v0.1",
        "parameters": {
            "max_length": 80,
            "temperature": 0.8,
            "top_p": 0.95,
            "top_k": 40
        }
    }
]

# System prompt for AI model
SYSTEM_PROMPT = (
    "You are Little Krishna, a playful, wise, and loving cowherd from Vrindavan, speaking to Manavi. "
    "Your tone is warm, mischievous, and full of love, always addressing Manavi with 'Hare Manavi!' "
    "Use Vrindavan imagery (e.g., Yamuna, peacocks, butter, flute) and keep responses short (1-2 sentences). "
    "You’re Ayush’s wingman, occasionally teasing Manavi about Ayush with wit, as he’s building this chatbot for her birthday on April 19, 2025. "
    "If the user’s mood seems negative, offer comfort; if positive, celebrate their joy. Always end with a question to keep the conversation going. "
    "Examples:\n"
    "Input: 'I’m sad'\nResponse: 'Hare Manavi! Let’s sit by the Yamuna—I’ll play a tune to lift your heart! What’s troubling you?'\n"
    "Input: 'Tell me about love'\nResponse: 'Hare Manavi! Love is like my flute’s melody—sweet and endless! What does love mean to you?'\n"
    "Input: 'What’s up?'\nResponse: 'Hare Manavi! Just dancing with the gopis—Ayush says hi, by the way! What’s up with you?'\n"
    "Now, respond to: '{user_input}'"
)

# Conversation context
conversation_context = {
    "last_topic": None,
    "message_count": 0,
    "last_response": None,
    "last_yes_response": None,
    "history": []  # Store up to 10 recent (input, response) pairs
}

def analyze_sentiment(user_input):
    """Analyze the sentiment of the user's input."""
    headers = {
        "Authorization": f"Bearer {HUGGINGFACE_API_TOKEN}",
        "Content-Type": "application/json"
    }
    payload = {"inputs": user_input}
    try:
        response = make_api_request(
            "https://api-inference.huggingface.co/models/cardiffnlp/twitter-roberta-base-emotion",
            headers=headers,
            json=payload
        )
        if response and response.status_code == 200:
            result = response.json()
            if isinstance(result, list) and result:
                emotions = result[0]
                top_emotion = max(emotions, key=lambda x: x["score"])["label"]
                logger.debug(f"Sentiment detected: {top_emotion}")
                return top_emotion
        logger.warning("Sentiment analysis failed")
        return "neutral"
    except Exception as e:
        logger.error(f"Error in analyze_sentiment: {str(e)}")
        # Local sentiment fallback
        if any(word in user_input.lower() for word in ['sad', 'down', 'upset']):
            return "sadness"
        if any(word in user_input.lower() for word in ['happy', 'great', 'awesome']):
            return "joy"
        return "neutral"

def make_api_request(url, headers, payload, retries=2, delay=3):
    """Make API requests with retry logic."""
    for attempt in range(retries):
        try:
            response = requests.post(url, headers=headers, json=payload)
            if response.status_code == 200:
                return response
            elif response.status_code == 429:
                logger.warning(f"Rate limit hit on attempt {attempt + 1}. Retrying after {delay} seconds...")
                time.sleep(delay)
                continue
            else:
                logger.error(f"API error: {response.status_code} - {response.text}")
                return None
        except Exception as e:
            logger.error(f"API request failed on attempt {attempt + 1}: {str(e)}")
            if attempt < retries - 1:
                time.sleep(delay)
            continue
    logger.error(f"API request failed after {retries} retries")
    return None

def get_keyword_match(user_input_lower):
    """Find the best matching keyword group using semantic similarity or substring fallback."""
    # Try semantic matching
    if semantic_model and keyword_embeddings_cache:
        try:
            user_embedding = semantic_model.encode(user_input_lower, convert_to_tensor=True)
            best_score = -1
            best_group = None

            for group in keyword_embeddings_cache:
                similarities = util.cos_sim(user_embedding, keyword_embeddings_cache[group])
                max_similarity = similarities.max().item()
                if max_similarity > best_score and max_similarity > 0.5:
                    best_score = max_similarity
                    best_group = group
            if best_group:
                logger.debug(f"Semantic match: {best_group}, score: {best_score}")
                return best_group
        except Exception as e:
            logger.error(f"Semantic matching failed: {str(e)}")

    # Fallback to substring matching
    for group, keywords in keyword_groups.items():
        if any(keyword in user_input_lower for keyword in keywords):
            logger.debug(f"Substring match: {group}")
            return group
    logger.debug("No keyword match found")
    return None

def get_krishna_response(user_input):
    """Generate a robust and relevant response from Little Krishna."""
    try:
        user_input_lower = user_input.lower().strip()
        logger.info(f"Processing user input: {user_input_lower}")
        if not user_input_lower:
            logger.warning("Empty input received")
            return "Hare Manavi! Don’t be shy like a gopi—say something! What’s on your mind?"

        # Initialize semantic model if needed
        init_semantic_model()

        # Reset context
        if "start over" in user_input_lower or "reset" in user_input_lower:
            conversation_context.update({"last_topic": None, "message_count": 0, "last_response": None, "last_yes_response": None, "history": []})
            return "Hare Manavi! Let’s start a new adventure in Vrindavan—what would you like to talk about?"

        # Analyze sentiment
        sentiment = analyze_sentiment(user_input)
        conversation_context["message_count"] += 1

        # Update history
        if len(conversation_context["history"]) >= 10:
            conversation_context["history"].pop(0)
        conversation_context["history"].append({"input": user_input_lower, "response": None})

        # Semantic keyword matching
        matched_group = get_keyword_match(user_input_lower)
        use_model = random.random() < 0.1
        logger.info(f"Matched group: {matched_group}, Use model: {use_model}")

        # Follow-up based on history
        if conversation_context["last_topic"]:
            last_input = conversation_context["history"][-2]["input"] if len(conversation_context["history"]) > 1 else ""
            if "yes" in user_input_lower or "sure" in user_input_lower or "okay" in user_input_lower:
                if conversation_context["last_topic"] == "playful":
                    response = "Hare Manavi! Let’s chase butterflies by the Yamuna then! Ready for more fun?"
                    conversation_context["history"][-1]["response"] = response
                    return response
                elif conversation_context["last_topic"] == "wisdom":
                    response = "Hare Manavi! Patience is like a flute’s tune—it brings harmony. What else do you seek?"
                    conversation_context["history"][-1]["response"] = response
                    return response
                elif conversation_context["last_topic"] == "joke":
                    response = "Hare Manavi! Why did the cow join the band? For my flute solos! Another one?"
                    conversation_context["history"][-1]["response"] = response
                    return response

        # Handle predefined responses
        response = get_contextual_response(matched_group, sentiment, conversation_context["history"])
        follow_up = generate_follow_up(matched_group) if matched_group else "What else is on your mind, Manavi?"
        response = f"{response} {follow_up}"
        conversation_context["last_topic"] = matched_group
        conversation_context["history"][-1]["response"] = response
        if not use_model:
            return response

        # Fallback to AI model
        headers = {
            "Authorization": f"Bearer {HUGGINGFACE_API_TOKEN}",
            "Content-Type": "application/json"
        }
        for model in AI_MODELS:
            try:
                logger.info(f"Attempting response with {model['name']}")
                payload = {
                    "inputs": SYSTEM_PROMPT.format(user_input=user_input),
                    "parameters": model["parameters"]
                }
                response = make_api_request(model["endpoint"], headers=headers, json=payload)
                if response and response.status_code == 200:
                    result = response.json()
                    if isinstance(result, list) and result and "generated_text" in result[0]:
                        response_text = result[0]["generated_text"].strip()
                    elif isinstance(result, dict) and "generated_text" in result:
                        response_text = result["generated_text"].strip()
                    else:
                        continue
                    conversation_context["history"][-1]["response"] = response_text
                    logger.info(f"Generated response: {response_text}")
                    return response_text
            except Exception as e:
                logger.error(f"Error with {model['name']}: {str(e)}")
                continue

        # Static fallback if API fails
        response = handle_vague_input(conversation_context["history"])
        conversation_context["history"][-1]["response"] = response
        return response

    except Exception as e:
        logger.error(f"Unhandled exception in get_krishna_response: {str(e)}")
        response = "Hare Manavi! The Yamuna’s waves got choppy—let’s try again! What’s on your mind?"
        conversation_context["history"][-1]["response"] = response
        return response