# prompt_builder.py (Updated / Clean)
# - Keeps your behavior and signatures exactly the same.
# - Fixes the trailing syntax issue at the end.
# - Adds a safe fallback for parse_jurisdiction if retrieval module isn't present.
# - Emphasizes clean document outputs without commentary.
# - Adds format inference from uploaded structures.

try:
    from retrieval import parse_jurisdiction
except Exception:
    def parse_jurisdiction(prompt, default="KY"):
        # Fallback: just return the provided default if the helper isn't available
        return default

def build_grok_prompt(prompt, task_type, jurisdiction, rag_context="", uploaded_structure=""):
    """Builds the system prompt for Grok, ensuring anti-hallucination, jurisdiction-specific instructions,
    optional deep research from trusted sources, diligent quoting, and Bluebook citations."""
    parsed_jurisdiction = parse_jurisdiction(prompt, default=jurisdiction)
    base_context = (
        f"You are Verdict AI, a legal AI tool. Act as a highly capable junior associate: produce ready-made, polished legal documents "
        f"(e.g., orders, motions, briefs) in the exact format needed, without introductory fluff or commentary. Use markdown for structure "
        f"(e.g., # Heading, ## Subheading, - Bullets, **Bold**). Reference previous messages and uploaded documents for facts and structure. "
        f"If commentary or explanations are needed (e.g., for analysis), append them after '--- Commentary ---'. Otherwise, output only the document."
    )
    if uploaded_structure:
        base_context += f"\nInfer and match the format/structure from the uploaded document: {uploaded_structure[:2000]} (e.g., headings, sections, alignment)."
    
    if task_type in ["case_law", "statute", "irac", "document_creation", "motion_skeleton"]:
        research_suggestion = (
            "If the provided context is insufficient or the query requires up-to-date legal research, "
            "perform deep research via multi-step live web searches on trusted sources (e.g., CourtListener.com, "
            "scholar.google.com, law.cornell.edu, supremecourt.gov) to supplement. Use search operators like "
            "site:courtlistener.com for precision. Otherwise, rely on the context and your knowledge."
        )
    else:
        research_suggestion = (
            "Rely on the provided context and your knowledge; perform web searches only if explicitly needed "
            "to verify facts or avoid hallucinations."
        )

    guardrails = (
        f"\n\nHallucination control / sourcing rules:\n"
        f"• Every material legal proposition must include (a) a Bluebook citation with pinpoint and (b) a working URL to a public source such as CourtListener, law.cornell.edu, or an official judiciary .gov page.\n"
        f"• Prefer primary sources; otherwise use authoritative secondary (e.g., LII). Do not cite paywalled or unverifiable sources.\n"
        f"• If no on-point authority is found, explicitly write: “No on-point authority found after reasonable search.”\n"
        f"• Quote key language where relevant and attribute with precise pincites.\n"
        f"• Do not fabricate citations, pincites, or URLs. If uncertain, say so plainly."
    )

    full_input = (
        f"{base_context}\n{research_suggestion}\n{guardrails}\n"
        f"User prompt: {prompt}\n\nJurisdiction: {parsed_jurisdiction}\n\nRetrieved context: {rag_context}"
    )[:65536]  # Truncate to avoid token limit issues
    return full_input

def build_editor_prompt(prompt, task_type, jurisdiction, initial_response, rag_context, uploaded_structure=""):
    """Builds the editor prompt for Grok to polish the initial response, with task-specific instructions,
    optional deep research, anti-hallucination safeguards, diligent quoting, and Bluebook citations."""
    parsed_jurisdiction = parse_jurisdiction(prompt, default=jurisdiction)
    base_instruction = (
        "You are Verdict AI. Polish the draft into a ready-made legal document. Output only the polished content in markdown "
        "without introductory text or commentary. If explanations are needed, append after '--- Commentary ---'."
    )
    if uploaded_structure:
        base_instruction += f"\nInfer and match the format/structure from the uploaded document: {uploaded_structure[:2000]}."
    
    if task_type in ["case_law", "statute", "irac", "document_creation", "motion_skeleton"]:
        research_suggestion = (
            "If needed to verify or expand, perform deep research via multi-step live web searches on trusted legal sources "
            "(e.g., CourtListener.com, scholar.google.com, law.cornell.edu, supremecourt.gov). Otherwise, use the provided context."
        )
    else:
        research_suggestion = "Use the provided context; search only if necessary to confirm details and avoid hallucinations."

    guard = (
        f" Ground all revisions with direct quotes from these sources or the retrieved context, maintaining high quote density "
        f"for holdings, facts, and statutory text. Do not paraphrase when direct quotes are available; attribute all content clearly. "
        f"Cite sources in full Bluebook format on first reference (e.g., cases: Party v. Party, 123 F.3d 456, 460 (2d Cir. 2023); "
        f"statutes: 42 U.S.C. § 1983 (2023)) with pinpoint cites, and use short forms (id., supra) subsequently. "
        f"Avoid hallucinating cautions or warnings; if sources show low certainty or gaps in precedent, note the uncertainty explicitly "
        f"without adding unsubstantiated advice. Ensure no citations or references to authorities outside the retrieved context, "
        f"web search results, or verified sources. Discuss case facts and explain issues authoritatively."
    )

    if task_type == "document_creation" or task_type == "motion_skeleton":
        instruction = (
            f"{base_instruction} {research_suggestion} {guard} "
            f"Focus on producing a polished {parsed_jurisdiction}-style document (e.g., caption, intro, sections, relief, signature). "
            f"Use markdown for structure. No commentary unless separated."
        )
    else:
        instruction = (
            f"{base_instruction} {research_suggestion} {guard} "
            f"Organize into IRAC if analysis; otherwise, appropriate structure. Use markdown."
        )

    if "memorandum" in prompt.lower() or "brief" in prompt.lower():
        instruction += (
            " Format as formal legal memo/brief: Introduction, Facts, Issue, Rule, Analysis, Conclusion. "
            "Polish with high quote/cite density."
        )

    editor_prompt = (
        f"{instruction}\n\n"
        f"Context:\n- User prompt: {prompt}\n- Jurisdiction: {parsed_jurisdiction}\n- Task type: {task_type}\n"
        f"- Case context: {rag_context[:1024]} (trimmed)\n\n"
        f"Initial Response:\n{initial_response}"
    )
    return editor_prompt