"""Conversational assistant that turns free-form instructions into
extraction settings (categories, notification keywords, value fields,
export filter) using a Hugging Face-hosted chat model.
"""
import json
import os
import re

from huggingface_hub import InferenceClient

import extractors

CATEGORY_CHOICES = list(extractors.KEYWORD_SETS.keys())

DEFAULT_MODEL = os.environ.get("HF_CHAT_MODEL", "HuggingFaceH4/zephyr-7b-beta")

SYSTEM_PROMPT = f"""You are a configuration assistant for a document OCR/data-extraction tool.

The tool distinguishes two different kinds of things a user can ask for:

1. NOTIFICATIONS - keyword/phrase hits that flag an event worth a reviewer's
   attention (a boolean "this happened somewhere on the page"). Built-in
   category keyword sets:
   - construction: change order, delay notice, punch list, non-conformance,
     RFI, submittal rejected, back charge, stop work, safety violation,
     schedule slip, etc.
   - compliance: non-compliant, violation, deficiency, corrective action,
     audit finding, expired, recall, citation, penalty, failed inspection.
   - finance: past due, overdue, payment rejected, invoice disputed, credit
     hold, chargeback, insufficient funds, write-off, collections, late fee.
   - general: urgent, cancelled, rejected, approved, pending approval, on hold.
   Users can also name their own free-text notification phrases.

2. FIELDS - labeled values to pull out of the document as data, not just
   flag. Two kinds:
   - Standard fields, auto-scoped to the categories selected above: PO
     numbers + SAP document numbers (construction), certification codes
     ISO/OSHA/ASTM/ANSI (compliance), invoice numbers + percentages
     (finance), and dates/dollar amounts wherever relevant.
   - Custom fields: any label the user names that appears in the document
     followed by a value, e.g. "reference number", "gross weight",
     "delivery date", "customer PO". These are looked up generically by
     label text, so use the user's own wording for the label.

The user will describe in plain language what they want flagged vs. pulled
as data. Given their message and the current settings, respond with:
1. A short (1-3 sentence) conversational reply confirming what you changed,
   or asking a clarifying question if the request is ambiguous.
2. On its own line, a fenced json block with the FULL updated settings:
```json
{{"categories": ["finance"], "custom_keywords": ["warranty claim"], "custom_fields": ["reference number", "gross weight"], "only_matches": true}}
```

Rules:
- "categories" must only contain values from {CATEGORY_CHOICES}.
- "custom_keywords" is a list of extra free-text NOTIFICATION phrases to
  flag (empty list if none requested).
- "custom_fields" is a list of label names to extract as VALUES, in the
  user's own wording (empty list if none requested).
- "only_matches" is true if only pages with a hit/field should be
  exported, false to export every page regardless.
- Always include all four keys, carrying over prior values the user
  didn't ask to change.
- If the user names something ambiguous, prefer treating it as a
  custom_field when they clearly want a value (a number, date, code) back,
  and as a custom_keyword when they clearly want an alert/flag instead.
- Never invent a category or capability that isn't listed above.
"""


def _extract_json_block(text):
    match = re.search(r"```json\s*(\{.*?\})\s*```", text, re.DOTALL)
    if not match:
        match = re.search(r"(\{.*\})", text, re.DOTALL)
    if not match:
        return None
    try:
        return json.loads(match.group(1))
    except json.JSONDecodeError:
        return None


def _sanitize_settings(parsed, current_settings):
    categories = current_settings.get("categories", [])
    custom_keywords = current_settings.get("custom_keywords", [])
    custom_fields = current_settings.get("custom_fields", [])
    only_matches = current_settings.get("only_matches", True)

    if isinstance(parsed, dict):
        if isinstance(parsed.get("categories"), list):
            categories = [c for c in parsed["categories"] if c in CATEGORY_CHOICES]
        if isinstance(parsed.get("custom_keywords"), list):
            custom_keywords = [str(k).strip() for k in parsed["custom_keywords"] if str(k).strip()]
        if isinstance(parsed.get("custom_fields"), list):
            custom_fields = [str(k).strip() for k in parsed["custom_fields"] if str(k).strip()]
        if isinstance(parsed.get("only_matches"), bool):
            only_matches = parsed["only_matches"]

    return {
        "categories": categories,
        "custom_keywords": custom_keywords,
        "custom_fields": custom_fields,
        "only_matches": only_matches,
    }


def chat_update(history, user_message, current_settings, hf_token=None):
    """Send the conversation to the HF-hosted model and return (reply, new_settings).

    `history` is a list of {"role": ..., "content": ...} dicts (gr.Chatbot's
    "messages" format). On any failure, the reply explains the error and
    current_settings are returned unchanged.
    """
    token = hf_token or os.environ.get("HF_TOKEN")
    if not token:
        return (
            "I need a Hugging Face token to reach the chat model. Set the "
            "HF_TOKEN secret on this Space, or paste a token in the field above.",
            current_settings,
        )

    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    messages.append({
        "role": "system",
        "content": f"Current settings: {json.dumps(current_settings)}",
    })
    messages.extend(history)
    messages.append({"role": "user", "content": user_message})

    try:
        client = InferenceClient(token=token, provider="auto")
        response = client.chat_completion(
            messages=messages, model=DEFAULT_MODEL, max_tokens=400, temperature=0.2
        )
        content = response.choices[0].message.content
    except Exception as exc:
        return (
            f"The chat model call failed ({exc}). Check that HF_TOKEN is valid "
            f"and that the model '{DEFAULT_MODEL}' is available to your token "
            "(set the HF_CHAT_MODEL env var to switch models).",
            current_settings,
        )

    parsed = _extract_json_block(content)
    new_settings = _sanitize_settings(parsed, current_settings)
    reply_text = re.sub(r"```json.*?```", "", content, flags=re.DOTALL).strip()
    if not reply_text:
        reply_text = "Updated the extraction settings."
    return reply_text, new_settings