import typer
from rich import print
from groq import Groq

# Replace with your Groq API key
API_KEY = "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"

client = Groq(api_key=API_KEY)

app = typer.Typer()


# AI function
def ask_ai(prompt):
    try:
        response = client.chat.completions.create(
            messages=[
                {"role": "user", "content": prompt}
            ],
            model="llama-3.1-8b-instant"
        )

        return response.choices[0].message.content

    except Exception as e:
        return f"Error contacting AI API: {e}"


# Ask command
@app.command()
def ask(question: str):
    """
    Ask AI assistant
    """
    print("[cyan]Thinking...[/cyan]")
    answer = ask_ai(question)
    print(f"[green]{answer}[/green]")


# Translate command
@app.command()
def translate(text: str, language: str):
    """
    Translate text to another language
    """
    prompt = f"Translate this to {language}: {text}"
    print("[cyan]Translating...[/cyan]")
    result = ask_ai(prompt)
    print(f"[green]{result}[/green]")


# Code generation command
@app.command()
def code(prompt: str):
    """
    Generate programming code
    """
    full_prompt = f"Write clean and correct code for: {prompt}"
    print("[cyan]Generating code...[/cyan]")
    result = ask_ai(full_prompt)
    print(f"[green]{result}[/green]")


# Explain code command
@app.command()
def explain(code: str):
    """
    Explain programming code
    """
    prompt = f"Explain this code:\n{code}"
    result = ask_ai(prompt)
    print(f"[green]{result}[/green]")


# Chat mode
@app.command()
def chat():
    """
    Start interactive AI chat
    """
    print("[bold green]Sumedh AI Chat started (type 'exit' to quit)[/bold green]")

    while True:
        msg = input("You: ")

        if msg.lower() == "exit":
            break

        response = ask_ai(msg)
        print(f"[cyan]AI:[/cyan] {response}")


# Main
if __name__ == "__main__":
    app()