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

# Enable detailed logging
logging.basicConfig(level=logging.INFO)

# Hugging Face Inference Client
client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")

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

### Function to analyze issues and provide solutions
def analyze_issues(issue_text: str, model_name: str, severity: str = None, programming_language: str = None) -> str:
    """
    Analyze issues and provide solutions.
    Args:
    issue_text (str): The issue text.
    model_name (str): The model name.
    severity (str, optional): The severity of the issue. Defaults to None.
    programming_language (str, optional): The programming language. Defaults to None.
    Returns:
    str: The analyzed issue and solution.
    """
    logging.info(f"Analyzing issue: {issue_text} with model: {model_name}")
    prompt = f"""Issue: {issue_text}
Severity: {severity}
Programming Language: {programming_language}
Please provide a comprehensive resolution in the following format:
## Problem Summary:
(Concise summary of the issue)
## Root Cause Analysis:
(Possible reasons for the issue)
## Solution Options:
1. **Option 1:** (Description)
   - Pros: (Advantages)
   - Cons: (Disadvantages)
2. **Option 2:** (Description)
   - Pros: (Advantages)
   - Cons: (Disadvantages)
## Recommended Solution:
(The best solution with justification)
## Implementation Steps:
1. (Step 1)
2. (Step 2)
3. (Step 3)
## Verification Steps:
1. (Step 1)
2. (Step 2)
"""
    try:
        nlp = pipeline("text-generation", model=model_name, max_length=1000)  # Increase max_length
        logging.info(f"Pipeline created with model: {model_name}")
        result = nlp(prompt)
        logging.info(f"Model output: {result}")
        return result[0]['generated_text']
    except Exception as e:
        logging.error(f"Error analyzing issue with model {model_name}: {e}")
        return f"Error analyzing issue with model {model_name}: {e}"

### Function to find related issues
def find_related_issues(issue_text: str, issues: list) -> list:
    """
    Find related issues.
    Args:
    issue_text (str): The issue text.
    issues (list): The list of issues.
    Returns:
    list: The list of related issues.
    """
    logging.info(f"Finding related issues for: {issue_text}")
    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))
    related_issues = sorted(related_issues, key=lambda x: x[1], reverse=True)
    logging.info(f"Found related issues: {related_issues}")
    return related_issues[:3]  # Return top 3 most similar issues

### Function to fetch GitHub issues
def fetch_github_issues(github_api_token: str, github_username: str, github_repository: str) -> list:
    """
    Fetch GitHub issues.
    Args:
    github_api_token (str): The GitHub API token.
    github_username (str): The GitHub username.
    github_repository (str): The GitHub repository.
    Returns:
    list: The list of GitHub issues.
    """
    logging.info(f"Fetching GitHub issues for: {github_username}/{github_repository}")
    url = f"https://api.github.com/repos/{github_username}/{github_repository}/issues"
    headers = {
        "Authorization": f"Bearer {github_api_token}",
        "Accept": "application/vnd.github.v3+json"
    }
    response = requests.get(url, headers=headers)
    if response.status_code == 200:
        issues = response.json()
        logging.info(f"Fetched issues: {issues}")
        return issues
    else:
        logging.error(f"Error fetching issues: {response.status_code}")
        raise Exception(f"Error fetching issues: {response.status_code}")

### Function to handle chat responses
def respond(
    command: str,
    history: list[tuple[str, str]],
    system_message: str,
    max_tokens: int,
    temperature: float,
    top_p: float,
    github_api_token: str,
    github_username: str,
    github_repository: str,
    selected_model: str,
    severity: str,
    programming_language: str,
    *args  # Add *args here
) -> str:
    """
    Handle chat responses.
    Args:
    command (str): The command.
    history (list[tuple[str, str]]): The chat history.
    system_message (str): The system message.
    max_tokens (int): The maximum number of tokens.
    temperature (float): The temperature.
    top_p (float): The top-p value.
    github_api_token (str): The GitHub API token.
    github_username (str): The GitHub username.
    github_repository (str): The GitHub repository.
    selected_model (str): The selected model.
    severity (str): The severity.
    programming_language (str): The programming language.
    Returns:
    str: The chat response.
    """
    global GITHUB_API_TOKEN
    GITHUB_API_TOKEN = github_api_token
    global issues
    issues = []

    messages = [{"role": "system", "content": system_message}]
    logging.info(f"System message: {system_message}")

    for user_msg, assistant_msg in history:
        if user_msg:
            messages.append({"role": "user", "content": user_msg})
            logging.info(f"User message: {user_msg}")
        if assistant_msg:
            messages.append({"role": "assistant", "content": assistant_msg})
            logging.info(f"Assistant message: {assistant_msg}")

    logging.info(f"Command received: {command}")

    if command == "/github":
    if not github_api_token:
        yield "Please enter your GitHub API token first. <https://github.com/settings/tokens>"
    else:
        try:
            issues = fetch_github_issues(github_api_token, github_username, github_repository)
            issue_list = "\n".join([f"{i+1}. {issue['title']}" for i, issue in enumerate(issues)])
            yield f"Available GitHub Issues:\n{issue_list}\n\nEnter the issue number to analyze:"
        except Exception as e:
            logging.error(f"Error fetching GitHub issues: {e}")
            yield f"Error fetching GitHub issues: {e}"

elif command == "/help":
    help_message = f"""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"""
    yield help_message  # Yield the pre-formatted help message

elif command.isdigit() and issues:
    try:
        issue_number = int(command) - 1
        issue = issues[issue_number]
        issue_text = issue['title'] + "\n\n" + issue['body']
        resolution = analyze_issues(issue_text, selected_model, severity, programming_language)

        # Find and display related issues
        related_issues = find_related_issues(issue_text, issues)
        related_issue_text = "\n".join(
            [f"- {issue['title']} (Similarity: {similarity:.2f})" for issue, similarity in related_issues]
        )

        yield f"Resolution for Issue '{issue['title']}':\n{resolution}\n\nRelated Issues:\n{related_issue_text}"
    except Exception as e:
        logging.error(f"Error analyzing issue: {e}")
        yield f"Error analyzing issue: {e}"

elif command.startswith("/generate_code"):
    # Extract the code description from the command
    code_description = command.replace("/generate_code", "").strip()
    if not code_description:
        yield "Please provide a description of the code you want to generate."
    else:
        prompt = f"Generate code for the following: {code_description}\nProgramming Language: {programming_language}"
        try:
            generated_code = analyze_issues(prompt, selected_model)
            code_output = f"<pre>{programming_language}\n{generated_code}</pre>"
            yield code_output  # Yield the formatted string
        except Exception as e:
            logging.error(f"Error generating code: {e}")
            yield f"Error generating code: {e}"

elif command.startswith("/explain_concept"):
    concept = command.replace("/explain_concept", "").strip()
    if not concept:
        yield "Please provide a concept to explain."
    else:
        prompt = f"Explain the concept of {concept} in detail."
        try:
            explanation = analyze_issues(prompt, selected_model)
            yield f"<pre>{explanation}</pre>"
        except Exception as e:
            logging.error(f"Error explaining concept: {e}")
            yield f"Error explaining concept: {e}"

elif command.startswith("/write_documentation"):
    topic = command.replace("/write_documentation", "").strip()
    if not topic:
        yield "Please provide a topic for documentation."
    else:
        prompt = f"Write comprehensive documentation for the following topic: {topic}"
        try:
            documentation = analyze_issues(prompt, selected_model)
            yield f"<pre>{documentation}</pre>"
        except Exception as e:
            logging.error(f"Error writing documentation: {e}")
            yield f"Error writing documentation: {e}"

elif command.startswith("/translate_code"):
    parts = command.replace("/translate_code", "").strip().split(" to ")
    if len(parts) != 2:
        yield "Invalid command format. Use: /translate_code [code] to [target language]"
    else:
        code, target_language = parts
        prompt = f"Translate the following code to {target_language}:\n\n


{code}\n

try:
translated_code = analyze_issues(prompt, selected_model)
code_output = f"<pre>{target_language}\n{translated_code}</pre>"
yield code_output # Yield the formatted string
except Exception as e:
logging.error(f"Error translating code: {e}")
yield f"Error translating code: {e}"

else:
messages.append({"role": "user", "content": command})
logging.info(f"User message: {command}")

response = ""
try:
    for message in client.chat_completion(
        messages,
        max_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        top_p=top_p,
    ):
        logging.info(f"Received message from chat completion: {message}")
        token = message.choices[0].delta.content
        response += token
        yield response
except Exception as e:
    logging.error(f"Error during chat completion: {e}")
    yield f"An error occurred: {e}"

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")

    # Define system_message here, after github_username and github_repository are defined
    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 Selection Dropdown
    model_dropdown = gr.Dropdown(
        choices=[
            "Xenova/gpt-4o",
            "acecalisto3/InstructiPhi",
            "DevShubham/Codellama-13B-Instruct-GGUF",
            "ricardo-larosa/SWE_Lite_dev-CodeLlama-34b",
            "DevsDoCode/Gemma-2b-Code-Instruct-Finetune-v0.1",
            "google/flan-t5-xxl",
            "facebook/bart-large-cnn",
            "microsoft/CodeBERT-base",
            "Salesforce/codegen-350M-mono",
            "bigcode/starcoder" 
        ],
        label="Select Model for Issue Resolution",
        value="microsoft/CodeBERT-base"
    )

    # Severity Dropdown
    severity_dropdown = gr.Dropdown(
        choices=["Critical", "Major", "Minor", "Trivial"],
        label="Severity",
        value=None  # Default to no severity selected
    )

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

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

    chatbot = gr.ChatInterface(
        respond,
        additional_inputs=[
            command_dropdown,  # Use command_dropdown instead of a regular message input
            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=1.1,
                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, show_header=False, debug=True)