import os
import sys
import io
from dotenv import load_dotenv
import google.generativeai as genai
from openai import OpenAI
import gradio as gr

# Load environment variables
load_dotenv()
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")

# Configure APIs
genai.configure(api_key=GEMINI_API_KEY)
openai = OpenAI(api_key=OPENAI_API_KEY)

# Model names
GEMINI_MODEL = "gemini-2.0-flash"  # Gemini model
OPENAI_MODEL = "gpt-4o-mini"  # OpenAI model (you can use "gpt-3.5-turbo" if needed)

def create_system_prompt():
    """Create the system prompt for the model."""
    system_message = "You are an assistant that adds high-quality docstrings and comments to Python code. "
    system_message += "Your task is to analyze the code and add appropriate documentation following PEP 257 conventions. "
    system_message += "Add function/class/module docstrings and inline comments where helpful. "
    system_message += "Don't change the functionality of the code, only add documentation. "
    system_message += "Use descriptive docstrings that explain parameters, return values, and exceptions raised."
    return system_message

def create_user_prompt(code):
    """Create the user prompt with the code to be documented."""
    user_prompt = "Please add appropriate docstrings and comments to the following Python code. "
    user_prompt += "Follow PEP 257 style for docstrings. Add comments only where they add value. "
    user_prompt += "Don't change the functionality of the code.\n\n"
    user_prompt += code
    return user_prompt

def document_code_with_gemini(code):
    """Generate documentation for the provided code using Gemini."""
    system_prompt = create_system_prompt()
    user_prompt = create_user_prompt(code)
    
    # Use Gemini to generate documentation
    model = genai.GenerativeModel(GEMINI_MODEL)
    response = model.generate_content(
        f"{system_prompt}\n\n{user_prompt}",
        generation_config=genai.GenerationConfig(
            temperature=0.7,
            max_output_tokens=2048,
        ),
    )
    
    # Clean up the response
    documented_code = response.text.replace("```python", "").replace("```", "").strip()
    return documented_code

def document_code_with_openai(code):
    """Generate documentation for the provided code using OpenAI."""
    system_prompt = create_system_prompt()
    user_prompt = create_user_prompt(code)
    
    # Use OpenAI to generate documentation
    response = openai.chat.completions.create(
        model=OPENAI_MODEL,
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ],
        temperature=0.7,
        max_tokens=2048,
    )
    
    # Clean up the response
    documented_code = response.choices[0].message.content.replace("```python", "").replace("```", "").replace("python", "")
    return documented_code

def document_code(code, api_choice):
    """Generate documentation using the selected API."""
    if api_choice == "Gemini":
        return document_code_with_gemini(code)
    elif api_choice == "OpenAI":
        return document_code_with_openai(code)
    else:
        return "Invalid API choice."

def execute_python(code):
    """Execute Python code and return the output."""
    try:
        output = io.StringIO()
        sys.stdout = output
        exec(code)
        return output.getvalue()
    except Exception as e:
        return f"Error: {str(e)}"
    finally:
        sys.stdout = sys.__stdout__

def create_interface():
    """Create the Gradio interface for the documentation tool."""
    sample_code = """
def fibonacci(n):
    fib_series = []
    a, b = 0, 1
    for _ in range(n):
        fib_series.append(a)
        a, b = b, a + b
    return fib_series

n = 10  
fib_series = fibonacci(n)
print("Fibonacci series up to", n, "terms:", fib_series)
"""

    css = """
    .original {background-color: #2d2d2d; color: #f8f8f2;}
    .documented {background-color: #1e1e1e; color: #f8f8f2;}

    /* Custom button styles */
    #add-documentation-btn {background-color: #007bff; color: white; font-weight: bold;} /* Blue for Add Documentation */
    #run-original-btn {background-color: #28a745; color: white; font-weight: bold;} /* Green for Run Original Code */
    #run-documented-btn {background-color: #6f42c1; color: white; font-weight: bold;} /* Purple for Run Documented Code */

    /* Ensure all buttons have white bold text */
    button {color: white; font-weight: bold;}
    """

    with gr.Blocks(css=css) as ui:
        gr.Markdown("## Automatic Python Documentation Generator")
        
        with gr.Row():
            input_code = gr.Textbox(label="Original Python code:", value=sample_code, lines=15, elem_classes=["original"])
            documented_code = gr.Textbox(label="Documented Python code:", lines=15, elem_classes=["documented"])
        
        with gr.Row():
            api_choice = gr.Radio(["Gemini", "OpenAI"], label="Select API", value="Gemini")
            document_btn = gr.Button("Add Documentation", elem_id="add-documentation-btn")
        
        with gr.Row():
            run_original = gr.Button("Run Original Code", elem_id="run-original-btn")
            run_documented = gr.Button("Run Documented Code", elem_id="run-documented-btn")
        
        with gr.Row():
            original_output = gr.TextArea(label="Original code output:")
            documented_output = gr.TextArea(label="Documented code output:")
        
        document_btn.click(document_code, inputs=[input_code, api_choice], outputs=[documented_code])
        run_original.click(execute_python, inputs=[input_code], outputs=[original_output])
        run_documented.click(execute_python, inputs=[documented_code], outputs=[documented_output])
        
    return ui

def main():
    """Main function to launch the interface."""
    ui = create_interface()
    ui.launch(inbrowser=True)

if __name__ == "__main__":
    main()