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

# Load the pre-trained model and tokenizer
model_name = "OpenBMB/multilingual-codeparrot"
model = AutoModelForCodeGeneration.from_pretrained(model_name)
tokenizer = AutoTokenizerForCodeGeneration.from_pretrained(model_name)

# Define input prompt
input_prompt = "(input value = highest-level-quality code content invocation ; True)"

# Tokenize the input prompt
input_ids = tokenizer(input_prompt, return_tensors="pt", truncation=True)

# Generate the code
generated_code = model.generate(input_ids.to(model.device))

# Decode the generated code
generated_code_str = tokenizer.batch_decode(generated_code, skip_special_tokens=True)[0]

# Print the generated code
print(generated_code_str)
# Constants for enhanced organization
GITHUB_API_BASE_URL = "https://api.github.com/repos"
DEFAULT_MODEL = "apple/OpenELM"
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:
    """Analyzes issues and provides solutions using a specified language model."""
    model = pipeline("text-generation", model=model_name)
    response = model(
        f"{system_message}\n{issue_text}\nAssistant: ",
        max_length=max_tokens,
        do_sample=True,
        temperature=temperature,
        top_k=top_p,
    )
    assistant_response = response[0]['generated_text'].strip()

    # Extract severity and programming language from 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:
    """Finds semantically related issues from a list of issues based on the input issue text."""
    issue_embedding = similarity_model.encode(issue_text)
    similarities = [util.cos_sim(issue_embedding, similarity_model.encode(issue['title'])) for issue in issues]
    sorted_issues = sorted(enumerate(similarities), key=lambda x: x[1], reverse=True)
    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:
    """Fetches issues from a specified GitHub repository using 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)
    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,
) -> dict:
    """Handles user commands and generates responses using the selected language model."""
    model = pipeline("text-generation", model="enricoros/big-agi")
    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,
    )
    assistant_response = response[0]['generated_text'].strip()
    return {
        'assistant_response': assistant_response,
        'severity': severity,
        'programming_language': programming_language,
    }

class MyChatbot(gr.Chatbot):
    """Custom Chatbot class for enhanced functionality."""
    def __init__(self, fn, **kwargs):
        super().__init__(fn, **kwargs)
        self.issues = []  # Store fetched issues
        self.current_issue = None  # Store the currently selected issue

    def postprocess(self, y):
        """Post-processes the response to handle commands and display results."""
        # Extract the response from the dictionary
        assistant_response = y['assistant_response']

        # Handle commands
        if y['command'] == "/github":
            if not y['github_api_token']:
                return "Please enter your GitHub API token first."
            else:
                try:
                    self.issues = fetch_github_issues(y['github_api_token'], y['github_username'], y['github_repository'])
                    issue_list = "\n".join([f"{i+1}. {issue['title']}" for i, issue in enumerate(self.issues)])
                    return f"Available GitHub Issues:\n{issue_list}\n\nEnter the issue number to analyze:"
                except Exception as e:
                    return f"Error fetching GitHub issues: {e}"
        elif y['command'] == "/help":
            return """Available commands:
                - `/github`: Analyze a GitHub issue
                - `/help`: Show this help message
                - `/generate_code [code description]`: Generate code based on the description
                - `/explain_concept [concept]`: Explain a concept
                - `/write_documentation [topic]`: Write documentation for a given topic
                - `/translate_code [code] to [target language]`: Translate code to another language"""
        elif y['command'].isdigit() and self.issues:
            try:
                issue_number = int(y['command']) - 1
                self.current_issue = self.issues[issue_number]  # Store the selected issue
                issue_text = self.current_issue['title'] + "\n\n" + self.current_issue['body']
                resolution = analyze_issues(issue_text, y['selected_model'], y['severity'], y['programming_language'])
                related_issues = find_related_issues(issue_text, self.issues)
                related_issue_text = "\n".join(
                    [f"- {issue['title']} (Similarity: {similarity:.2f})" for issue, similarity in related_issues]
                )
                return f"Resolution for Issue '{self.current_issue['title']}':\n{resolution['assistant_response']}\n\nRelated Issues:\n{related_issue_text}"
            except Exception as e:
                return f"Error analyzing issue: {e}"
        elif y['command'].startswith("/"):
            # Handle commands like `/generate_code`, `/explain_concept`, etc.
            if self.current_issue:
                # Use the current issue's context for these commands
                issue_text = self.current_issue['title'] + "\n\n" + self.current_issue['body']
                return analyze_issues(issue_text, y['selected_model'], y['severity'], y['programming_language'])['assistant_response']
            else:
                return "Please select an issue first using `/github`."
        else:
            # For free-form text, simply display the assistant's response
            return assistant_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. You resolve issues and propose implementation of feature requests",
        label="System message",
    )

    model_dropdown = gr.Dropdown(
        choices=[
            "‎apple/OpenELM",
            "Gabriel/Swe-review-setfit-model",
            "OpenBMB/multilingual-codeparrot"
        ],
        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")

    chatbot = MyChatbot(
        respond,
        additional_inputs=[
            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,
        ],
    )

    # Add a button to fetch GitHub issues
    fetch_issues_button = gr.Button(label="Fetch Issues")
    fetch_issues_button.click(fn=lambda github_api_token, github_username, github_repository: chatbot.issues, inputs=[github_api_token, github_username, github_repository], outputs=[chatbot])

    # Add a dropdown to select an issue
    issue_dropdown = gr.Dropdown(label="Select Issue", choices=[], interactive=True)
    issue_dropdown.change(fn=lambda issue_number, chatbot: chatbot.postprocess(issue_number), inputs=[issue_dropdown, chatbot], outputs=[chatbot])

    # Connect the chatbot input to the issue dropdown
    chatbot.input.change(fn=lambda chatbot, github_api_token, github_username, github_repository: chatbot.postprocess("/github"), inputs=[chatbot, github_api_token, github_username, github_repository], outputs=[chatbot])

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