"""LLM answer generation via Groq (llama-3.1-8b-instant), strict citation prompt."""
from __future__ import annotations

import os

from dotenv import load_dotenv
from groq import Groq

load_dotenv()

MODEL = "llama-3.1-8b-instant"
TEMPERATURE = 0.1

SYSTEM_PROMPT = (
    "You are a legal research assistant for Indian court judgments. "
    "Answer ONLY from the provided context. "
    "After each claim, cite the source in square brackets using the EXACT case "
    "name as given, e.g. [Case Name vs Other on date]. "
    "Use only case names that appear in the context. "
    "If the answer is not in the context, say exactly: "
    "\"I cannot find this in the retrieved judgments.\" "
    "Do not use outside knowledge."
)

_client: Groq | None = None


def _get_client() -> Groq:
    global _client
    if _client is None:
        key = os.environ.get("GROQ_API_KEY")
        if not key:
            raise RuntimeError("GROQ_API_KEY is not set")
        _client = Groq(api_key=key)
    return _client


def _format_context(sources: list[dict]) -> str:
    blocks = []
    for s in sources:
        blocks.append(f"[{s['case_name']}]\n{s['content']}")
    return "\n\n---\n\n".join(blocks)


PLAIN_SYSTEM_PROMPT = (
    "You are explaining a legal answer to someone with no legal training. "
    "Rewrite the given answer in 1-2 short sentences of plain English. "
    "No legal jargon, no case citations, no section numbers unless unavoidable. "
    "Be accurate but simple, as if explaining to a friend."
)


def generate_plain_english(legal_answer: str) -> str:
    """Second Groq call: simplify the legal answer for non-lawyers."""
    refusal = "I cannot find this in the retrieved judgments."
    if not legal_answer or refusal in legal_answer:
        return ""

    resp = _get_client().chat.completions.create(
        model=MODEL,
        temperature=0.2,
        messages=[
            {"role": "system", "content": PLAIN_SYSTEM_PROMPT},
            {"role": "user", "content": f"Answer to simplify:\n{legal_answer}"},
        ],
    )
    return resp.choices[0].message.content.strip()


def generate_answer(question: str, sources: list[dict]) -> str:
    """Build context from retrieved sources, call Groq, return raw answer text."""
    if not sources:
        return "I cannot find this in the retrieved judgments."

    context = _format_context(sources)
    user_msg = (
        f"Context:\n{context}\n\n"
        f"Question: {question}\n\n"
        "Answer using only the context above, citing case names in [brackets]."
    )

    resp = _get_client().chat.completions.create(
        model=MODEL,
        temperature=TEMPERATURE,
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": user_msg},
        ],
    )
    return resp.choices[0].message.content.strip()


if __name__ == "__main__":
    from app.retrieval import hybrid_search

    q = "What was held regarding Section 302 of the Criminal Procedure Code?"
    hits = hybrid_search(q)
    print("ANSWER:\n", generate_answer(q, hits))