import gradio as gr
import google.generativeai as genai
import paramiko
import io
import re

dependency_file_patterns = "|".join([
    "package\.json", "requirements\.txt", "Gemfile", "pom\.xml", "build\.gradle",
    "composer\.json", "Cargo\.toml", "go\.mod", "project\.clj", "mix\.exs",
    "build\.sbt", "Package\.swift", ".*\.csproj", "packages\.config", "yarn\.lock",
    "package-lock\.json", "Pipfile", "poetry\.lock", "environment\.yml",
    "pubspec\.yaml", "deps\.edn", "rebar\.config", "bower\.json", "Podfile", "Cartfile"
])

def fetch_github_file(github_url, ssh_private_key):
    try:
        # Parse the GitHub SSH URL
        parts = github_url.split(':')
        owner_repo = parts[1].split('.git')[0]
        owner, repo = owner_repo.split('/')
        
        # Set up SSH client
        ssh = paramiko.SSHClient()
        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())

        # Load the private key
        private_key = paramiko.RSAKey.from_private_key(io.StringIO(ssh_private_key))

        # Connect to GitHub using SSH
        ssh.connect('github.com', username='git', pkey=private_key)

        # Construct the git command to fetch the file content
        git_command = f"git ls-tree -r --name-only HEAD | grep -E '({dependency_file_patterns})'"

        # Execute the command
        stdin, stdout, stderr = ssh.exec_command(git_command)
        
        # Read the file list
        file_list = stdout.read().decode('utf-8').splitlines()

        file_contents = {}
        for file_path in file_list:
            cat_command = f"git show HEAD:{file_path}"
            stdin, stdout, stderr = ssh.exec_command(cat_command)
            file_contents[file_path] = stdout.read().decode('utf-8')

        # Close the SSH connection
        ssh.close()

        if not file_contents:
            return "Error: No dependency files found."

        return file_contents

    except paramiko.AuthenticationException:
        return "Error: Authentication failed. Please check your SSH key."
    except paramiko.SSHException as ssh_err:
        return f"SSH Error: {str(ssh_err)}"
    except Exception as e:
        return f"Error accessing GitHub: {str(e)}"

def process_with_gemini(file_contents, gemini_api_key):
    genai.configure(api_key=gemini_api_key)
    model = genai.GenerativeModel('gemini-pro')

    combined_content = "\n\n".join([f"File: {name}\n{content}" for name, content in file_contents.items()])

    prompt = f"""
    Analyze the following dependency files:

    {combined_content}

    Please provide:
    1. A list of dependencies and their versions for each file
    2. The licenses associated with each dependency (if available)
    3. A summary of the project based on these dependencies
    4. Any potential license conflicts or considerations
    5. Suggestions for best practices in open-source license management for this project
    """

    response = model.generate_content(prompt)
    return response.text

def process_input(file, github_url, ssh_private_key, gemini_api_key):
    if file is not None and github_url:
        return "Error: Please either upload a file OR provide a GitHub URL, not both."
    
    if file is not None:
        file_content = {file.name: file.read().decode('utf-8')}
    elif github_url and ssh_private_key:
        # ... (GitHub URL handling)
        file_content = fetch_github_file(github_url, ssh_private_key)
        if isinstance(file_content, str) and file_content.startswith("Error:"):
            return file_content
    else:
        return "Error: Please either upload a file OR provide both GitHub URL and SSH Private Key."
    
    try:
        # Process the file content with Gemini
        analysis = process_with_gemini(file_content, gemini_api_key)
        return analysis
    except Exception as e:
        return f"Error processing the file: {str(e)}"

iface = gr.Interface(
    fn=process_input,
    inputs=[
        gr.File(label="Upload dependency file (e.g., requirements.txt, package.json, Gemfile)"),
        gr.Textbox(label="GitHub Repository URL (optional, use SSH format)"),
        gr.Textbox(label="SSH Private Key (required if using GitHub URL)", type="password"),
        gr.Textbox(label="Gemini API Key", type="password"),
    ],
    outputs=gr.Textbox(label="License Information and Analysis"),
    title="Open Source License Extractor",
    description="Upload a dependency file or provide a GitHub repository URL to analyze open-source licenses.",
)

if __name__ == "__main__":
    iface.launch()