import os
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from groq import Groq
from dotenv import load_dotenv

# Load environment variables from a .env file
load_dotenv()  # reads .env in project root by default

API_KEY = os.getenv("GROQ_API_KEY")
if not API_KEY:
    raise RuntimeError("Please set GROQ_API_KEY in your .env file")

# Initialize Groq client once
client = Groq(api_key=API_KEY)

def translate_text(
    text: str,
    target_language: str,
    model: str = "llama-3.3-70b-versatile",
    temperature: float = 0.7,
    stream: bool = False
) -> str:
    """
    Sends a translation prompt to the LLM and returns the translated text.
    """
    messages = [
        {"role": "system", "content": "You are a helpful, accurate translator."},
        {
            "role": "user",
            "content": (
                f"Translate the following text into {target_language}:\n\n"
                f"{text}"
            )
        }
    ]

    completion = client.chat.completions.create(
        model=model,
        messages=messages,
        temperature=temperature,
        max_completion_tokens=1024,
        top_p=1,
        stream=stream,
        stop=None,
    )

    if not stream:
        return completion.choices[0].message.content.strip()
    else:
        result = []
        for chunk in completion:
            delta = chunk.choices[0].delta.content
            if delta:
                result.append(delta)
        return "".join(result).strip()

# FastAPI setup
app = FastAPI(title="Groq Translator API")

class TranslateRequest(BaseModel):
    text: str
    target_language: str
    stream: bool = False

class TranslateResponse(BaseModel):
    translation: str

@app.post("/translate", response_model=TranslateResponse)
async def translate_endpoint(req: TranslateRequest):
    """
    POST /translate
    {
      "text": "Hello, world!",
      "target_language": "French",
      "stream": false
    }
    =>
    {
      "translation": "Bonjour, le monde !"
    }
    """
    try:
        translated = translate_text(
            text=req.text,
            target_language=req.target_language,
            stream=req.stream
        )
        return TranslateResponse(translation=translated)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# To run:
# 1. pip install fastapi uvicorn groq python-dotenv
# 2. create a .env file in the project root with:
#      GROQ_API_KEY=your_real_api_key_here
# 3. uvicorn main:app --reload --port 8000