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

# 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"  # Robust default model selection
MAX_RELATED_ISSUES = 3 

# 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:
    """
    Analyzes issues and provides comprehensive solutions using a specified language model.

    Args:
        issue_text (str): The detailed description of the issue.
        model_name (str): The name of the language model to use for analysis.
        severity (str, optional): The severity level of the issue (e.g., "Critical", "Major"). Defaults to None.
        programming_language (str, optional): The programming language relevant to the issue. Defaults to None.

    Returns:
        str: A structured analysis of the issue, including problem summary, root cause, solution options, 
             recommended solution, implementation steps, and verification steps.
    """
    logging.info(f"Analyzing issue: '{issue_text}' with model: '{model_name}'")

    # Construct a detailed prompt for the language model
    prompt = f"""Issue: {issue_text}
Severity: {severity if severity else "Not specified"}
Programming Language: {programming_language if programming_language else "Not specified"}

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)
        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. Please try again later." # User-friendly error message

### Function to find related issues
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.

    Args:
        issue_text (str): The text of the issue for which to find related issues.
        issues (list): A list of issues, where each issue is a dictionary with at least a 'title' key.

    Returns:
        list: A list of tuples, each containing a related issue (dictionary) and its similarity score to the input issue.
              The list is sorted by similarity in descending order, limited to MAX_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)[:MAX_RELATED_ISSUES]
    logging.info(f"Found related issues: {related_issues}")
    return related_issues

### Function to fetch GitHub 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.

    Args:
        github_api_token (str): The user's GitHub API token for authentication.
        github_username (str): The username of the GitHub account owning the repository.
        github_repository (str): The name of the GitHub repository.

    Returns:
        list: A list of dictionaries, where each dictionary represents a GitHub issue with its details.

    Raises:
        Exception: If there's an error fetching issues from the GitHub API.
    """
    logging.info(f"Fetching GitHub issues for: '{github_username}/{github_repository}'")
    url = f"{GITHUB_API_BASE_URL}/{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: Status Code {response.status_code}, Response: {response.text}")
        raise Exception(f"Error fetching GitHub issues. Please check your credentials and repository.")

### 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,
) -> str:
    """
    Handles user commands and generates responses using the selected language model.

    Args:
        command (str): The user's command or input.
        history (list[tuple[str, str]]): The conversation history as a list of (user, assistant) message tuples.
        system_message (str): The initial system message defining the chatbot's persona.
        max_tokens (int): The maximum number of tokens allowed in the generated response.
        temperature (float): The temperature parameter controlling the randomness of the generated text.
        top_p (float): The top-p parameter for nucleus sampling, controlling the diversity of the generated text.
        github_api_token (str): The user's GitHub API token for authentication.
        github_username (str): The username of the GitHub account owning the repository.
        github_repository (str): The name of the GitHub repository.
        selected_model (str): The name of the language model selected for generating responses.
        severity (str): The severity level of the issue (e.g., "Critical", "Major").
        programming_language (str): The programming language relevant to the issue.

    Returns:
        str: The chatbot's response, generated by the selected language model or based on the command.
    """

    global issues 
    issues = [] # Reset issues for each interaction

    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:
            return "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)])
                return 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}")
                return "Error fetching GitHub issues. Please check your credentials and repository."

    elif command == "/help":
        help_message = """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"""
        return 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]
            )

            return 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}")
            return "Error analyzing issue. Please try again later."

    elif command.startswith("/generate_code"):
        # Extract the code description from the command
        code_description = command.replace("/generate_code", "").strip()
        if not code_description:
            return "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>{generated_code}</pre>"
                return code_output
            except Exception as e:
                logging.error(f"Error generating code: {e}")
                return "Error generating code. Please try again later."

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

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

    elif command.startswith("/translate_code"):
        try:
            code_and_language = command.replace("/translate_code", "").strip().split(" to ")
            code = code_and_language[0]
            target_language = code_and_language[1]
            prompt = f"Translate the following code to {target_language}:\n\n{code}"
            try:
                translated_code = analyze_issues(prompt, selected_model)
                code_output = f"<pre>{translated_code}</pre>"
                return code_output
            except Exception as e:
                logging.error(f"Error translating code: {e}")
                return "Error translating code. Please try again later."
        except Exception as e:
            logging.error(f"Error parsing translate_code command: {e}")
            return "Error parsing translate_code command. Please check the syntax."

    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 # Use yield for streaming responses
        except Exception as e:
            logging.error(f"Error during chat completion: {e}")
            return "An error occurred. Please try again later."

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",
            "microsoft/CodeBERT-base"
        ],
        label="Select Model for Issue Resolution",
        value=DEFAULT_MODEL # Set a default model
    )

    # 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,
            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, debug=False)