import os
from openai import OpenAI
import gradio as gr
import anthropic
from dotenv import load_dotenv

load_dotenv(override=True)

openai_api_key = os.getenv('OPENAI_API_KEY')
anthropic_api_key = os.getenv('ANTHROPIC_API_KEY')

if openai_api_key:
    print(f"OpenAI API Key exists and begins {openai_api_key[:8]}")
else:
    print("OpenAI API Key not set")
if anthropic_api_key:
    print(f"Anthropic API Key exists and begins {anthropic_api_key[:7]}")
else:
    print("Anthropic API Key not set")

MODEL = "gpt-4o-mini"
openai = OpenAI()
CLAUDE_MODEL='claude-3-7-sonnet-latest'
claude= anthropic.Anthropic()

system_message = """
You are a Python expert assistant.

Your job is to add a clean, detailed, and professional Python docstring to a given function.

The docstring should:
- Start with a one-line summary of what the function does.
- Include an `Args:` section listing all parameters with type and description.
- Include a `Returns:` section describing the return value with type.
- Follow proper indentation and formatting standards (PEP 257 style).
- Do not modify the function body.
- Do not explain or comment — just return the updated function with the docstring inserted.

Example:

Input:
def multiply(a, b):
    return a * b

Output:
def multiply(a, b):
    \"\"\"
    Multiply two numbers.

    Args:
        a (int): The first number.
        b (int): The second number.

    Returns:
        int: The product of a and b.
    \"\"\"
    return a * b
"""
system_message_unittest = """
You are a Python testing assistant.

Your job is to generate clean, correct, and complete unit tests for the given Python function(s) using the built-in `unittest` framework.

Your output should:
- Import the function if needed (assume it's in the same file)
- Use a test class that inherits from `unittest.TestCase`
- Include 2 to 4 relevant test cases using `.assertEqual`, `.assertTrue`, or `.assertFalse`
- Cover edge cases (e.g. empty inputs, zero, negative, large numbers)
- Not explain the code — just return the full Python code block with the unit tests

Example:

Input:
def is_even(n):
    return n % 2 == 0

Output:
import unittest

class TestIsEven(unittest.TestCase):
    def test_even(self):
        self.assertTrue(is_even(4))
        self.assertTrue(is_even(0))

    def test_odd(self):
        self.assertFalse(is_even(3))

if __name__ == "__main__":
    unittest.main()
"""

def userPrompt(code):
    userPrompt=code
    return userPrompt

def generator(code, mode):
    response = ""
    if mode == "docstring":
        response = openai.chat.completions.create(
            model=MODEL,
            messages=[
                {"role":"system","content":system_message},
                {"role":"user","content":userPrompt(code)}
            ]
        )
    elif mode == "unit_test":
        response = openai.chat.completions.create(
            model=MODEL,
            messages=[
                {"role":"system","content":system_message_unittest},
                {"role":"user","content":userPrompt(code)}
            ]
        )
    gpt_output = response.choices[0].message.content
    # Remove triple backticks and 'python' label
    if gpt_output.startswith("```python"):
        gpt_output = gpt_output.replace("```python", "").strip()
    if gpt_output.endswith("```"):
        gpt_output = gpt_output[:-3].strip()
        
    # Unescape \n and convert to real line breaks
    cleaned = gpt_output.encode().decode('unicode_escape')

    return cleaned.strip()

interface = gr.Interface(
    fn=generator,
    inputs=[
        gr.Textbox(lines=12, label="Paste your Python function here"),
        gr.Dropdown(choices=["docstring", "unit_test"], label="Select Mode", value="docstring")
    ],
    outputs=gr.Code(label="Output", language="python"),
    title="🧠 Python Docstring and Unit Test Generator",
    description="Paste your Python function and select the mode to generate a docstring or a unit test automatically."
)

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