import os
import json
from groq import AsyncGroq

groq_client = AsyncGroq(api_key=os.getenv("GROQ_API_KEY", ""))

async def extract_citations(text: str) -> dict:
    """
    Uses LLaMA3-70B zero-shot to extract and format citations from raw text.
    """
    # Only check the last 4000 characters where references usually are
    text_snippet = text[-4000:] if len(text) > 4000 else text
    
    prompt = f"""
    Tugas Anda adalah mengekstrak daftar pustaka/referensi dari teks di bawah ini.
    Format output HARUS berupa JSON object dengan key 'citations' yang berisi array of objects.
    Setiap object memiliki kunci:
    - 'author': nama pengarang
    - 'year': tahun publikasi
    - 'title': judul jurnal/buku
    - 'venue': nama jurnal atau penerbit (opsional)
    
    Jika tidak ada daftar pustaka yang ditemukan, kembalikan {{"citations": []}}.
    Hanya kembalikan JSON murni, tanpa penjelasan apapun.
    
    Teks:
    {text_snippet}
    """
    
    try:
        response = await groq_client.chat.completions.create(
            model="llama-3.3-70b-versatile",
            messages=[{"role": "user", "content": prompt}],
            temperature=0.1,
            max_tokens=1024,
            response_format={"type": "json_object"}
        )
        
        content = response.choices[0].message.content
        data = json.loads(content)
        return {"success": True, "data": data.get("citations", [])}
    except Exception as e:
        return {"success": False, "error": str(e)}