import gradio as gr
import git
import os
from transformers import pipeline, AutoTokenizer, AutoModelForQuestionAnswering


def clone_repo(repo_url):
    local_path = "repo_clone"
    git.Repo.clone_from(repo_url, local_path)
    return local_path


def process_repo(repo_url, option):
    if option == "Pre-trained":
        qa_pipeline = pipeline('question-answering')
    else:
        model_path = "./model"
        tokenizer = AutoTokenizer.from_pretrained(model_path)
        model = AutoModelForQuestionAnswering.from_pretrained(model_path)
        qa_pipeline = pipeline('question-answering', model=model, tokenizer=tokenizer)

    repo_path = clone_repo(repo_url)
    result = {}
    for root, dirs, files in os.walk(repo_path):
        for file in files:
            file_path = os.path.join(root, file)
            with open(file_path, 'r', encoding="utf-8") as f:
                text = f.read()
                summary = text[:50] + "..." if len(text) > 50 else text
                keywords = qa_pipeline(summary)['answer']
                result[file_path] = {"summary": summary, "text": text, "keywords": keywords}
    return result


def qa_chatbot(repo_dict, question):
    all_text = ""
    for file in repo_dict.values():
        all_text += file['summary'] + " " + file['text'] + " "
    answer = qa_pipeline({'context': all_text, 'question': question})['answer']
    return answer


input_repo = gr.inputs.Textbox(label="Enter Git repository URL")
output_processed_repo = gr.outputs.Textbox(label="Processed Git repository")
output_qa_chatbot = gr.outputs.Textbox(label="Answer")

model_options = ["Pre-trained", "Fine-tuned"]
input_option = gr.inputs.Dropdown(choices=model_options, label="Choose a model option")

process_repo_interface = gr.Interface(fn=process_repo, inputs=[input_repo, input_option], outputs=output_processed_repo,
                                      title="Process Git Repository")

qa_chatbot_interface = gr.Interface(fn=qa_chatbot, inputs={"repo_dict": gr.inputs.Dictionary(
    key_type=gr.inputs.Textbox(label="File path"),
    value_type=gr.inputs.Dictionary(key_type=gr.inputs.Textbox(label="File content"),
                                    value_type=gr.inputs.Textbox(label="Keywords"))),
    "question": gr.inputs.Textbox(label="Question")},
    outputs=output_qa_chatbot,
    title="QA Chatbot")

process_repo_interface.launch()
qa_chatbot_interface.launch()