import json
from sentence_transformers import SentenceTransformer, util
from groq import Groq
from datetime import datetime
import os
import pandas as pd
from datasets import load_dataset, Dataset
from dotenv import load_dotenv
import random
import glob

# Load environment variables
load_dotenv()

# Initialize Groq client
groq_client = Groq(api_key=os.getenv("GROQ_API_KEY"))

# Load similarity model
similarity_model = SentenceTransformer('paraphrase-MiniLM-L6-v2')

# Config
HF_DATASET_REPO = "midrees2806/unmatched_queries"
HF_TOKEN = os.getenv("HF_TOKEN")

# Greeting list
GREETINGS = [
    "hi", "hello", "hey", "good morning", "good afternoon", "good evening",
    "assalam o alaikum", "salam", "aoa", "hi there",
    "hey there", "greetings"
]

# Fixed rephrased unmatched query responses
UNMATCHED_RESPONSES = [
    "Thank you for your query. We’ve forwarded it to our support team and it will be added soon. In the meantime, you can visit the University of Education official website or reach out via the contact details below.\n\n📞 +92-42-99262231-33\n✉️ info@ue.edu.pk\n🌐 https://ue.edu.pk",
    "We’ve noted your question and it’s in queue for inclusion. For now, please check the University of Education website or contact the administration directly.\n\n📞 +92-42-99262231-33\n✉️ info@ue.edu.pk\n🌐 https://ue.edu.pk",
    "Your query has been recorded. We’ll update the system with relevant information shortly. Meanwhile, you can visit UE's official site or reach out using the details below:\n\n📞 +92-42-99262231-33\n✉️ info@ue.edu.pk\n🌐 https://ue.edu.pk",
    "We appreciate your question. It has been forwarded for further processing. Until it’s available here, feel free to visit the official UE website or use the contact options:\n\n📞 +92-42-99262231-33\n✉️ info@ue.edu.pk\n🌐 https://ue.edu.pk"
]

# Load multiple JSON datasets
dataset = []
try:
    json_files = glob.glob('datasets/*.json')
    for file_path in json_files:
        with open(file_path, 'r', encoding='utf-8') as f:
            data = json.load(f)
            if isinstance(data, list):
                for item in data:
                    if isinstance(item, dict) and 'Question' in item and 'Answer' in item:
                        dataset.append(item)
                    else:
                        print(f"Invalid entry in {file_path}: {item}")
            else:
                print(f"File {file_path} does not contain a list.")
except Exception as e:
    print(f"Error loading datasets: {e}")

# Precompute embeddings
dataset_questions = [item.get("Question", "").lower().strip() for item in dataset]
dataset_answers = [item.get("Answer", "") for item in dataset]
dataset_embeddings = similarity_model.encode(dataset_questions, convert_to_tensor=True)

# Save unmatched queries to Hugging Face
def manage_unmatched_queries(query: str):
    try:
        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        try:
            ds = load_dataset(HF_DATASET_REPO, token=HF_TOKEN)
            df = ds["train"].to_pandas()
        except:
            df = pd.DataFrame(columns=["Query", "Timestamp", "Processed"])
        if query not in df["Query"].values:
            new_entry = {"Query": query, "Timestamp": timestamp, "Processed": False}
            df = pd.concat([df, pd.DataFrame([new_entry])], ignore_index=True)
            updated_ds = Dataset.from_pandas(df)
            updated_ds.push_to_hub(HF_DATASET_REPO, token=HF_TOKEN)
    except Exception as e:
        print(f"Failed to save query: {e}")

# Query Groq LLM
def query_groq_llm(prompt, model_name="llama3-70b-8192"):
    try:
        chat_completion = groq_client.chat.completions.create(
            messages=[{
                "role": "user",
                "content": prompt
            }],
            model=model_name,
            temperature=0.7,
            max_tokens=500
        )
        return chat_completion.choices[0].message.content.strip()
    except Exception as e:
        print(f"Error querying Groq API: {e}")
        return ""

# Main logic function to be called from Gradio
def get_best_answer(user_input):
    if not user_input.strip():
        return "Please enter a valid question."

    user_input_lower = user_input.lower().strip()

    if len(user_input_lower.split()) < 3 and not any(greet in user_input_lower for greet in GREETINGS):
        return "Please ask your question properly with at least 3 words."

    if any(keyword in user_input_lower for keyword in ["fee structure", "fees structure", "semester fees", "semester fee"]):
        return (
            "💰 For the most complete and up-to-date fee details for your program at the University of Education Lahore, please visit the official fee structure page.\n"
            "This webpage offers a detailed overview of the fee structure, providing you with essential information to support your academic journey at our institution.\n"
            "🔗 https://drive.google.com/file/d/1B30FKoP6GrkS9pQk10PWKCwcjco5E9Cc/view"
        )

    user_embedding = similarity_model.encode(user_input_lower, convert_to_tensor=True)
    similarities = util.pytorch_cos_sim(user_embedding, dataset_embeddings)[0]
    best_match_idx = similarities.argmax().item()
    best_score = similarities[best_match_idx].item()

    if best_score < 0.65:
        manage_unmatched_queries(user_input)
        return random.choice(UNMATCHED_RESPONSES)

    original_answer = dataset_answers[best_match_idx]
    prompt = f"""Name is UOE AI Assistant! You are an official assistant for the University of Education Lahore.
Rephrase the following official answer clearly and professionally. 
Use structured formatting (like headings, bullet points, or numbered lists) where appropriate.
DO NOT add any new or extra information. ONLY rephrase and improve the clarity and formatting of the original answer.
### Question:
{user_input}
### Original Answer:
{original_answer}
### Rephrased Answer:
"""

    llm_response = query_groq_llm(prompt)

    if llm_response:
        for marker in ["Improved Answer:", "Official Answer:", "Rephrased Answer:"]:
            if marker in llm_response:
                return llm_response.split(marker)[-1].strip()
        return llm_response
    else:
        return dataset_answers[best_match_idx]