import gradio as gr
from transformers import pipeline
from sentence_transformers import SentenceTransformer, util
import os
import requests
import logging

# Configure logging for detailed insights
logging.basicConfig(level=logging.INFO)

# Constants for enhanced organization
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) -> str:
    # Initialize the model
    model = pipeline("text-generation", model=model_name)

    # Generate a response
    response = model(
        f"{system_message}\n{issue_text}\nAssistant: ",
        max_length=max_tokens,
        do_sample=True,
        temperature=temperature,
        top_k=top_p,
    )

    # Extract the assistant's response
    assistant_response = response[0]['generated_text'].strip()

    # Analyze the response
    if "Severity" in assistant_response:
        severity = assistant_response.split(":")[1].strip()

    if "Programming Language" in assistant_response:
        programming_language = assistant_response.split(":")[1].strip()

    return {
        'assistant_response': assistant_response,
        'severity': severity,
        'programming_language': programming_language,
    }

def find_related_issues(issue_text: str, issues: list) -> list:
    # Calculate the similarity between the issue and other issues
    issue_embedding = similarity_model.encode(issue_text)
    similarities = [util.cos_sim(issue_embedding, similarity_model.encode(issue['title'])) for issue in issues]

    # Sort the issues by similarity
    sorted_issues = sorted(enumerate(similarities), key=lambda x: x[1], reverse=True)

    # Select the top related issues
    related_issues = [issues[i] for i, similarity in sorted_issues[:MAX_RELATED_ISSUES]]

    return related_issues

def fetch_github_issues(github_api_token: str, github_username: str, github_repository: str) -> list:
    # Fetch the issues from the GitHub API
    headers = {'Authorization': f'token {github_api_token}'}
    url = f"{GITHUB_API_BASE_URL}/{github_username}/{github_repository}/issues"
    response = requests.get(url, headers=headers)

    # Parse the JSON response
    issues = response.json()

    return issues

def respond(
    command,
    history,
    system_message,
    max_tokens,
    temperature,
    top_p,
    github_api_token,
    github_username,
    github_repository,
    selected_model,
    severity,
    programming_language,
    *args,
    **kwargs,
) -> str:
    # Initialize the model
    model = pipeline("text-generation", model=selected_model)

    # Generate a response
    response = model(
        f"{system_message}\n{command}\n{history}\n{github_username}/{github_repository}\n{severity}\n{programming_language}\nAssistant: ",
        max_length=max_tokens,
        do_sample=True,
        temperature=temperature,
        top_k=top_p,
    )

    # Extract the assistant's response
    assistant_response = response[0]['generated_text'].strip()

    return assistant_response
class MyChatbot(gr.ChatInterface):
    def __init__(self, fn, *args, **kwargs):
        super().__init__(fn, *args, **kwargs)

    def update_chat_history(self, message: str, is_user: bool) -> None:
        if is_user:
            self.history.append((message, None))
        else:
            self.history.append((None, message))

    def compute(self, *args, **kwargs):
        command = args[0]
        history = self.history
        system_message = self.additional_inputs["system_message"]
        max_tokens = self.additional_inputs["max_new_tokens"]
        temperature = self.additional_inputs["temperature"]
        top_p = self.additional_inputs["top_p"]
        github_api_token = self.additional_inputs["github_api_token"]
        github_username = self.additional_inputs["github_username"]
        github_repository = self.additional_inputs["github_repository"]
        selected_model = self.additional_inputs["model_dropdown"]
        severity = self.additional_inputs["severity_dropdown"]
        programming_language = self.additional_inputs["programming_language_textbox"]

        return self.respond(
            command,
            history,
            system_message,
            max_tokens,
            temperature,
            top_p,
            github_api_token,
            github_username,
            github_repository,
            selected_model,
            severity,
            programming_language,
        )

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. You resolve issues and propose implementation of feature requests",
        label="System message",
    )

    model_dropdown = gr.Dropdown(
        choices=[
            "microsoft/CodeBERT-base",
            "Salesforce/codegen-350M-mono",
        ],
        label="Select Model for Issue Resolution",
        value=DEFAULT_MODEL,
    )

    severity_dropdown = gr.Dropdown(
        choices=["Critical", "Major", "Minor", "Trivial"],
        label="Severity",
        value=None,
    )

    programming_language_textbox = gr.Textbox(label="Programming Language")

    command_dropdown = gr.Dropdown(
        choices=[
            "/github",
            "/help",
            "/generate_code",
            "/explain_concept",
            "/write_documentation",
            "/translate_code",
        ],
        label="Select Command",
    )

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

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