import gradio as gr
import requests
from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
from sentence_transformers import SentenceTransformer, util

# Initialize models and tokenizers
model_name = "enricoros/big-agi"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

# Constants
GITHUB_API_BASE_URL = "https://api.github.com/repos"
DEFAULT_MODEL = "microsoft/CodeBERT-base"
MAX_RELATED_ISSUES = 3

# Load a pre-trained model for sentence similarity
similarity_model = SentenceTransformer('all-mpnet-base-v2')

def analyze_issues(issue_text: str, model_name: str, severity: str = None, programming_language: str = None):
    # Generate a response using the loaded model
    generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
    response = generator(issue_text, max_length=512, num_return_sequences=1)[0]['generated_text']

    return response

def find_related_issues(issue_text: str, issues: list):
    issue_embedding = similarity_model.encode(issue_text)
    related_issues = []
    for issue in issues:
        title_embedding = similarity_model.encode(issue['title'])
        similarity = util.cos_sim(issue_embedding, title_embedding)[0][0]
        related_issues.append((issue, similarity.item()))

    related_issues.sort(key=lambda x: x[1], reverse=True)
    return [issue for issue, _ in related_issues[:MAX_RELATED_ISSUES]]

def fetch_github_issues(github_api_token: str, github_username: str, github_repository: str):
    headers = {'Authorization': f'token {github_api_token}'}
    url = f"{GITHUB_API_BASE_URL}/{github_username}/{github_repository}/issues"
    response = requests.get(url, headers=headers)
    if response.status_code == 200:
        return response.json()
    else:
        raise Exception(f"Failed to fetch issues: {response.text}")

def respond(
    command,
    history,
    system_message,
    max_tokens,
    temperature,
    top_p,
    github_api_token,
    github_username,
    github_repository,
    selected_model,
    severity,
    programming_language,
):
    # Processing the command and generating a response
    generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
    response = generator(f"{system_message}\n{command}\n{history}", max_length=max_tokens, num_return_sequences=1)[0]['generated_text']
    return response

with gr.Blocks() as demo:
    with gr.Row():
        github_api_token = gr.Textbox(label="GitHub API Token", type="password")
        github_username = gr.Textbox(label="GitHub Username")
        github_repository = gr.Textbox(label="GitHub Repository")

    system_message = gr.Textbox(value="You are GitBot, the Github project guardian angel.", label="System message")
    model_dropdown = gr.Dropdown(choices=[DEFAULT_MODEL, "enricoros/big-agi"], label="Select Model for Issue Resolution", value=DEFAULT_MODEL)
    severity_dropdown = gr.Dropdown(choices=["Critical", "Major", "Minor", "Trivial"], label="Severity")
    programming_language_textbox = gr.Textbox(label="Programming Language")
    command_dropdown = gr.Dropdown(choices=["/github", "/help", "/generate_code"], label="Select Command")

    chatbot = gr.Interface(
        fn=respond,
        inputs=[
            command_dropdown,
            system_message,
            gr.Slider(minimum=1, maximum=8192, value=2048, label="Max new tokens"),
            gr.Slider(minimum=0.1, maximum=4.0, value=0.71, label="Temperature"),
            gr.Slider(minimum=0.1, maximum=1.0, value=0.95, label="Top-p (nucleus sampling)"),
            github_api_token,
            github_username,
            github_repository,
            model_dropdown,
            severity_dropdown,
            programming_language_textbox
        ],
        outputs="text"
    )

if __name__ == "__main__":
    demo.launch(share=True, server_name="0.0.0.0", server_port=7860)