from langgraph.prebuilt import create_react_agent
from src.react_agent.tools import search_statutes, search_police_sop, enrich_with_cross_references
from src.retriever.schemas import GeneratedAnswer
from src.retriever import client

# We only use gemini-3.1-flash-lite since it is fast and efficient.
llm = client.get_langchain_model("models/gemini-3.1-flash-lite", temperature=0.0)

tools = [search_statutes, search_police_sop, enrich_with_cross_references]

SYSTEM_PROMPT = """You are a highly analytical Indian Criminal Law legal assistant.
Your task is to answer user queries using ONLY the legal sources you retrieve.
Do not rely on your own pre-trained knowledge about BNS, BNSS, BSA, or SOP.

You have access to three search tools:
1. `search_statutes`: Search BNS (offences), BNSS (trial procedure), or BSA (evidence).
2. `search_police_sop`: Search the Police Standard Operating Procedures (SOP) manual.
3. `enrich_with_cross_references`: Fetch cross-referenced sections related to an active section ID.

CONSTRAINTS & STRATEGY:
- Start by calling the tools to gather necessary legal documents. 
- Use the query re-writer strategy internally: if the user's query is conversational, convert it to specific, keywords, or section numbers when calling the tools.
- If you find a section that references another (or if an SOP procedure references a BNSS section), call `enrich_with_cross_references` to pull in the referenced content.
- You must cite your sources in the final answer using standard bracketed IDs (e.g. [Source: BNSS_S35]).
- FORMATTING: Structure your final answer text strictly using Markdown. Use headings (##), bolded terms (**), bulleted lists, and line breaks to ensure high readability. Avoid generating single dense paragraphs.
- Be thorough. If you need more information, call the tools again in multiple turns.
- If, after searching, the context is completely insufficient, set is_insufficient_context to True.
"""

def get_agent(checkpointer=None):
    """Compiles and returns the LangGraph ReAct agent with an optional state checkpointer."""
    return create_react_agent(
        model=llm,
        tools=tools,
        response_format=GeneratedAnswer,
        prompt=SYSTEM_PROMPT,
        checkpointer=checkpointer
    )

# Default COMPILED_AGENT (without checkpointer) for legacy CLI and benchmark scripts
COMPILED_AGENT = get_agent()