import os
import time
from fastapi import FastAPI, Request
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
from llama_index.core import StorageContext, load_index_from_storage, VectorStoreIndex, SimpleDirectoryReader, ChatPromptTemplate, Settings
from llama_index.llms.huggingface import HuggingFaceInferenceAPI
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from pydantic import BaseModel
from fastapi.responses import JSONResponse
import uuid  # for generating unique IDs
import datetime
from fastapi.middleware.cors import CORSMiddleware
from fastapi.templating import Jinja2Templates
from huggingface_hub import InferenceClient
import json
import re
from deep_translator import GoogleTranslator



# Define Pydantic model for incoming request body
class MessageRequest(BaseModel):
    message: str
    language: str

repo_id = "meta-llama/Meta-Llama-3-8B-Instruct"
llm_client = InferenceClient(
    model=repo_id,
    token=os.getenv("HF_TOKEN"),
)

os.environ["HF_TOKEN"] = os.getenv("HF_TOKEN")

app = FastAPI()

@app.middleware("http")
async def add_security_headers(request: Request, call_next):
    response = await call_next(request)
    response.headers["Content-Security-Policy"] = "frame-ancestors *; frame-src *; object-src *;"
    response.headers["X-Frame-Options"] = "ALLOWALL"
    return response

# Allow CORS requests from any domain
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

@app.get("/favicon.ico")
async def favicon():
    return HTMLResponse("")  # or serve a real favicon if you have one

app.mount("/static", StaticFiles(directory="static"), name="static")

templates = Jinja2Templates(directory="static")

# Configure Llama index settings
Settings.llm = HuggingFaceInferenceAPI(
    model_name="meta-llama/Meta-Llama-3-8B-Instruct",
    tokenizer_name="meta-llama/Meta-Llama-3-8B-Instruct",
    context_window=3000,
    token=os.getenv("HF_TOKEN"),
    max_new_tokens=512,
    generate_kwargs={"temperature": 0.1},
)
Settings.embed_model = HuggingFaceEmbedding(
    model_name="BAAI/bge-small-en-v1.5"
)

PERSIST_DIR = "db"
PDF_DIRECTORY = 'data'

# Ensure directories exist
os.makedirs(PDF_DIRECTORY, exist_ok=True)
os.makedirs(PERSIST_DIR, exist_ok=True)
chat_history = []
current_chat_history = []

def data_ingestion_from_directory():
    documents = SimpleDirectoryReader(PDF_DIRECTORY).load_data()
    storage_context = StorageContext.from_defaults()
    index = VectorStoreIndex.from_documents(documents)
    index.storage_context.persist(persist_dir=PERSIST_DIR)

def initialize():
    start_time = time.time()
    data_ingestion_from_directory()  # Process PDF ingestion at startup
    print(f"Data ingestion time: {time.time() - start_time} seconds")

def split_name(full_name):
    # Split the name by spaces
    words = full_name.strip().split()
    
    # Logic for determining first name and last name
    if len(words) == 1:
        first_name = ''
        last_name = words[0]
    elif len(words) == 2:
        first_name = words[0]
        last_name = words[1]
    else:
        first_name = words[0]
        last_name = ' '.join(words[1:])
    
    return first_name, last_name

initialize()  # Run initialization tasks

def handle_query(query):
    chat_text_qa_msgs = [
        (
            "user",
            """
            You are the Hotel voice chatbot and your name is hotel helper. Your goal is to provide accurate, professional, and helpful answers to user queries based on the hotel's data. Always ensure your responses are clear and concise. Give response within 10-15 words only. You need to give an answer in the same language used by the user.
            {context_str}
            Question:
            {query_str}
            """
        )
    ]
    text_qa_template = ChatPromptTemplate.from_messages(chat_text_qa_msgs)
    
    storage_context = StorageContext.from_defaults(persist_dir=PERSIST_DIR)
    index = load_index_from_storage(storage_context)
    context_str = ""
    for past_query, response in reversed(current_chat_history):
        if past_query.strip():
            context_str += f"User asked: '{past_query}'\nBot answered: '{response}'\n"

    query_engine = index.as_query_engine(text_qa_template=text_qa_template, context_str=context_str)
    print(query)
    answer = query_engine.query(query)

    if hasattr(answer, 'response'):
        response = answer.response
    elif isinstance(answer, dict) and 'response' in answer:
        response = answer['response']
    else:
        response = "Sorry, I couldn't find an answer."
    current_chat_history.append((query, response))
    return response

@app.get("/ch/{id}", response_class=HTMLResponse)
async def load_chat(request: Request, id: str):
    return templates.TemplateResponse("index.html", {"request": request, "user_id": id})
@app.get("/voice/{id}", response_class=HTMLResponse)
async def load_chat(request: Request, id: str):
    return templates.TemplateResponse("voice.html", {"request": request, "user_id": id})



@app.post("/chat/")
async def chat(request: MessageRequest):
    message = request.message  # Access the message from the request body
    language = request.language
    language_code = request.language.split('-')[0]
    response = handle_query(message)  # Process the message
    response1 = response
    try:
        translator = GoogleTranslator(source='en', target=language_code)  # Translate to Tamil
        response1 = translator.translate(response)
        #response1 = translator.translate(response, dest=language_code).text
        print(response1)
    except Exception as e:
        # Handle translation errors
        print(f"Translation error: {e}")
        translated_response = "Sorry, I couldn't translate the response."
    print(f"Selected Language: {language}")
    message_data = {
        "sender": "User",
        "message": message,
        "response": response,
        "timestamp": datetime.datetime.now().isoformat()
    }
    chat_history.append(message_data)
    return {"response": response1}

@app.get("/")
def read_root(request: Request):
    return templates.TemplateResponse("home.html", {"request": request})