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


# Define Pydantic model for incoming request body
class MessageRequest(BaseModel):
    message: str
repo_id = "meta-llama/Meta-Llama-3-8B-Instruct"
llm_client = InferenceClient(
    model=repo_id,
    token=os.getenv("HF_TOKEN"),
)
def summarize_conversation(inference_client: InferenceClient, history: list):
    # Construct the full prompt with history
    history_text = "\n".join([f"{entry['sender']}: {entry['message']}" for entry in history])
    full_prompt = f"{history_text}\n\nSummarize the conversation in three concise points only give me only Summarization in python list formate :\n"

    response = inference_client.post(
        json={
            "inputs": full_prompt,
            "parameters": {"max_new_tokens": 512},
            "task": "text-generation",
        },
    )
    
    # Decode the response
    generated_text = json.loads(response.decode())[0]["generated_text"]

    # Use regex to extract the list inside brackets
    matches = re.findall(r'\[(.*?)\]', generated_text)

    # If matches found, extract the content
    if matches:
        # Assuming we only want the first match, split by commas and strip whitespace
        list_items = matches[0].split(',')
        cleaned_list = [item.strip() for item in list_items]
        return cleaned_list
    else:
        return generated_text


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")

initialize()  # Run initialization tasks


def handle_query(query):
    chat_text_qa_msgs = [
        (
            "user",
            """
            You are the Clara Redfernstech chatbot. Your goal is to provide accurate, professional, and helpful answers to user queries based on the company's data. Always ensure your responses are clear and concise. Give response within 10-15 words only       
            {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)
    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})
# Route to save chat history
@app.post("/hist/")
async def save_chat_history(history: dict):
    # Check if 'history' is a key in the incoming dictionary
    user_id = history.get('userId')
    print(user_id)
    if 'history' in history and isinstance(history['history'], list):
        print("Received history:", history['history'])  # Debugging line
        cleaned_summary = summarize_conversation(llm_client, history['history'])
        print("Cleaned summary:", cleaned_summary)  # Debugging line
        return {"summary": cleaned_summary, "message": "Chat history saved"}
    else:
        return JSONResponse(status_code=400, content={"message": "Invalid history format"})
@app.post("/webhook")
async def receive_form_data(request: Request):
    form_data = await request.json()
    
    # Generate a unique ID (for tracking user)
    unique_id = str(uuid.uuid4())
    
    # Here you can do something with form_data like saving it to a database
    print("Received form data:", form_data)
    
    # Send back the unique id to the frontend
    return JSONResponse({"id": unique_id})

@app.post("/chat/")
async def chat(request: MessageRequest):
    message = request.message  # Access the message from the request body
    response = handle_query(message)  # Process the message
    message_data = {
        "sender": "User",
        "message": message,
        "response": response,
        "timestamp": datetime.datetime.now().isoformat()
    }
    chat_history.append(message_data)
    return {"response": response}
@app.get("/")
def read_root():
    return {"message": "Welcome to the API"}