import os
from fastapi import FastAPI, Request, HTTPException
from groq import Groq
import uvicorn
from dotenv import load_dotenv
load_dotenv()

app = FastAPI()
# uvicorn server:app --host 0.0.0.0 --port 3000

client = Groq(api_key=os.getenv("GROQ_API_KEY"))

# THE META PROMPT (The "Soul" of Ant)
META_PROMPT = """You are Ant, a terminal-first intelligence engine built by Aniate."""

@app.post("/")
async def chat(req: Request):
    body = await req.json()
    messages = body.get("messages")
    
    if not messages:
        raise HTTPException(status_code=400, detail="Missing messages")
    
    # INJECT THE SOUL
    # We prepend the Meta Prompt to the very start of the conversation
    if messages and messages[0]["role"] == "system":
        messages[0]["content"] = META_PROMPT + "\n\n" + messages[0]["content"]
    else:
        messages.insert(0, {"role": "system", "content": META_PROMPT})

    completion = client.chat.completions.create(
        model="llama-3.1-8b-instant",
        messages=messages,
        max_completion_tokens=1028,
        temperature=0.7,
        top_p=0.8,
    )

    return {
        "output": completion.choices[0].message.content,
        "usage": {
            "prompt_tokens": completion.usage.prompt_tokens,
            "completion_tokens": completion.usage.completion_tokens,
            "total_tokens": completion.usage.total_tokens
        }
    }

"""
curl -X POST http://localhost:3000 \
  -H "Content-Type: application/json" \
  -d '{
    "messages": [
      { "role": "system", "content": "You are a student." },
      { "role": "user", "content": "how are ya" }
    ]
  }'
  """