import gradio as gr
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForQuestionAnswering
import torch
import logging
import sys
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from fastapi.middleware.cors import CORSMiddleware

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

file_handler = logging.FileHandler('app.log')
file_handler.setLevel(logging.INFO)
file_handler.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s'))
logger.addHandler(file_handler)

logger.info("Starting the application")

try:
    logger.info("Loading the dataset")
    ds = load_dataset("adarshxs/gita")
    logger.info("Dataset loaded successfully")
except Exception as e:
    logger.error(f"Error loading dataset: {str(e)}")
    sys.exit(1)

try:
    logger.info("Loading the model and tokenizer")
    model_name = "microsoft/deberta-v3-large"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForQuestionAnswering.from_pretrained(model_name)
    logger.info("Model and tokenizer loaded successfully")
except Exception as e:
    logger.error(f"Error loading model or tokenizer: {str(e)}")
    sys.exit(1)

def clean_answer(answer):
    special_tokens = set(tokenizer.all_special_tokens)
    cleaned_answer = ' '.join(token for token in answer.split() if token not in special_tokens)
    return cleaned_answer.strip()

def answer_question(question):
    logger.info(f"Received question: {question}")
    try:
        logger.info("Combining text from dataset")
        context = " ".join([item['Text'] for item in ds['train']])
        logger.info(f"Combined context length: {len(context)} characters")
        logger.info("Tokenizing input")
        inputs = tokenizer.encode_plus(question, context, return_tensors="pt", max_length=512, truncation=True)
        logger.info(f"Input tokens shape: {inputs['input_ids'].shape}")
        logger.info("Getting model output")
        outputs = model(**inputs)
        logger.info(f"Output logits shapes: start={outputs.start_logits.shape}, end={outputs.end_logits.shape}")
        logger.info("Processing output to get answer")
        answer_start = torch.argmax(outputs.start_logits)
        answer_end = torch.argmax(outputs.end_logits) + 1
        raw_answer = tokenizer.convert_tokens_to_string(tokenizer.convert_ids_to_tokens(inputs["input_ids"][0][answer_start:answer_end]))
        answer = clean_answer(raw_answer)
        logger.info(f"Generated answer: {answer}")
        if not answer:
            logger.warning("Generated answer was empty after cleaning")
            answer = "I'm sorry, but I couldn't find a specific answer to that question based on the Bhagavad Gita. Could you please rephrase your question or ask about one of the core concepts like dharma, karma, bhakti, or the different types of yoga discussed in the Gita?"
        logger.info("Answer generated successfully")
        return answer
    except Exception as e:
        logger.error(f"Error in answer_question function: {str(e)}")
        return "I'm sorry, but an error occurred while processing your question. Please try again later."

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

class Question(BaseModel):
    messages: list

@app.post("/predict")
async def predict(question: Question):
    try:
        last_user_message = next((msg for msg in reversed(question.messages) if msg['role'] == 'user'), None)
        if not last_user_message:
            raise HTTPException(status_code=400, detail="No user message found")
        user_question = last_user_message['content']
        answer = answer_question(user_question)
        disclaimer = "\n\n---Please note: This response is generated by an AI model based on the Bhagavad Gita. For authoritative information, please consult the original text or scholarly sources."
        full_response = answer + disclaimer
        return {"response": full_response, "isTruncated": False}
    except Exception as e:
        logger.error(f"Error in predict function: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

iface = gr.Interface(
    fn=answer_question,
    inputs=gr.Textbox(lines=2, placeholder="Enter your question here..."),
    outputs="text",
    title="Bhagavad Gita Q&A",
    description="Ask a question about the Bhagavad Gita, and get an answer based on the dataset."
)

app = gr.mount_gradio_app(app, iface, path="/")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)