# task_processing.py (Updated for new VerdictAI modes)
# - Adds labels/headers for new modes introduced in app.py
# - Keeps lightweight, post-processing-only behavior
# - No redundant model calls (no extra Grok/GPT invocations)
# - Safe default passthrough if a mode/task_type is unrecognized

# Retained import in case you reference it elsewhere; not used here.
from gpt_helpers import ask_gpt41_mini  # noqa: F401


def _wrap(title: str, jurisdiction: str, content: str) -> str:
    """Minimal, consistent wrapper for returning labeled content."""
    jur = f" ({jurisdiction})" if jurisdiction else ""
    return f"{title}{jur}:\n\n{content}"


def process_task_response(task_type: str, grok_response: str, prompt: str, jurisdiction: str) -> str:
    """
    Post-process Grok output for display/download. This function is intentionally lightweight:
    - No additional retrieval or model calls
    - Just adds consistent headings or minimal formatting per task type
    - Safe passthrough for unknown types
    """
    tt = (task_type or "").lower()

    # Legacy/primary types
    if tt in ("case_law", "irac"):
        return _wrap("Case Law / IRAC Analysis", jurisdiction, grok_response)
    if tt == "statute":
        return _wrap("Statute Analysis", jurisdiction, grok_response)
    if tt == "document_creation":
        return _wrap("Final Document", jurisdiction, grok_response)
    if tt == "document_analysis":
        return _wrap("Document Analysis", jurisdiction, grok_response)
    if tt == "legal_strategy":
        return _wrap("Legal Strategy", jurisdiction, grok_response)

    # Existing misc types (kept for compatibility with any callers you might have)
    if tt == "analogical_reasoning":
        return _wrap("Analogical Reasoning", jurisdiction, grok_response)
    if tt in ("deposition_questions", "deposition_qs"):
        return _wrap("Deposition Questions", jurisdiction, grok_response)
    if tt == "verdict_predictor":
        return _wrap("Verdict Prediction", jurisdiction, grok_response)
    if tt == "damages_estimator":
        return _wrap("Damages Estimate", jurisdiction, grok_response)
    if tt == "pleading_validator":
        return _wrap("Pleading Validation", jurisdiction, grok_response)
    if tt == "legal_inbox":
        return _wrap("Legal Inbox Processing", jurisdiction, grok_response)
    if tt == "privilege_analyzer":
        return _wrap("Privilege/PII Analysis", jurisdiction, grok_response)
    if tt == "eli5":
        return _wrap("Plain-English Explanation", jurisdiction, grok_response)
    if tt == "workflow_template":
        return _wrap("Workflow Template", jurisdiction, grok_response)

    # New modes introduced in app.py
    if tt == "citation_check":
        return _wrap("Citation Check / Shepardizing-lite", jurisdiction, grok_response)
    if tt == "brief_builder":
        return _wrap("Brief Builder (IRAC Sections)", jurisdiction, grok_response)
    if tt in ("deposition_qa",):
        return _wrap("Deposition Q&A Generator", jurisdiction, grok_response)
    if tt == "checklist":
        return _wrap("Checklist & Red Flags", jurisdiction, grok_response)
    if tt == "comparative_law":
        return _wrap("Comparative Law", jurisdiction, grok_response)
    if tt == "client_summary":
        return _wrap("Client-Friendly Summary", jurisdiction, grok_response)
    if tt == "motion_skeleton":
        return _wrap("Motion Skeleton", jurisdiction, grok_response)
    if tt == "opposing_argument":
        return _wrap("Opposing Argument (Devil's Advocate)", jurisdiction, grok_response)
    if tt == "case_extractor":
        return _wrap("Case/Citation Extractor", jurisdiction, grok_response)
    if tt == "judge_style":
        return _wrap("Judge/County Style Adaptation", jurisdiction, grok_response)
    if tt == "general_qa":
        return _wrap("General Q&A", jurisdiction, grok_response)

    # Fallback passthrough
    return grok_response