"""Shared, offline-testable product recognition for the demo and evaluation."""
from __future__ import annotations
import re
import unicodedata
from collections import Counter, defaultdict

FORM_WORDS = {
    "tbl", "cps", "sol", "crm", "gel", "ung", "inj", "inf", "plv", "sus", "sir", "gtt", "spr",
    "flm", "obd", "eff", "ret", "pro", "por", "ijf", "dis", "lyo", "aer", "emp", "pst", "liq",
}


def normalize(text: str) -> str:
    """Lowercase, strip diacritics and punctuation (Slovak inflection is handled by prefix matching)."""
    text = unicodedata.normalize("NFKD", text)
    text = "".join(ch for ch in text if not unicodedata.combining(ch))
    return re.sub(r"[^a-z0-9]+", " ", text.lower()).strip()


def build_aliases(_products: list[dict]) -> tuple[dict, dict, dict]:
    """Map normalized names to products.

    Returns (aliases, by_prefix, by_url): aliases maps an alias to
    {"kind": "substance"|"product", "label": str, "products": [index, ...]}.
    """
    aliases: dict[str, dict] = {}
    substance_products: dict[str, list[int]] = defaultdict(list)
    brand_products: dict[str, list[int]] = defaultdict(list)

    for i, product in enumerate(_products):
        for substance in product["subs"]:
            key = normalize(substance)
            if len(key) >= 4:
                substance_products[key].append(i)
        tokens = normalize(product["name"]).split()
        brand = []
        for token in tokens:
            if any(ch.isdigit() for ch in token) or token in FORM_WORDS:
                if brand:
                    break
                continue  # leading "3M", "0 9 %" ... are not the brand
            brand.append(token)
        if brand and len(brand[0]) >= 3:
            brand_products[brand[0]].append(i)
            if len(brand) >= 2:
                brand_products[" ".join(brand[:2])].append(i)

    for key, ids in substance_products.items():
        label = Counter(s for i in ids for s in _products[i]["subs"] if normalize(s) == key).most_common(1)[0][0]
        aliases[key] = {"kind": "substance", "label": label, "products": ids}
    for key, ids in brand_products.items():
        if key in aliases:
            continue
        # A word shared by many unrelated products (e.g. "vitamin") is not a brand.
        if len({tuple(_products[i]["subs"]) for i in ids}) > 6:
            continue
        aliases[key] = {"kind": "product", "label": key.title(), "products": ids}

    by_prefix: dict[str, list[str]] = defaultdict(list)
    for key in aliases:
        by_prefix[key[:4]].append(key)
    by_url = {p["url"]: p for p in _products if p["url"]}
    return aliases, by_prefix, by_url


def detect_drugs(question: str, aliases: dict, by_prefix: dict) -> list[dict]:
    """Find products and active substances mentioned in the question.

    Slovak nouns are inflected (warfarínom, ibuprofénu), so a question word
    matches an alias it starts with, allowing a short ending.
    """
    tokens = normalize(question).split()
    found, used = {}, set()
    for n in (2, 1):
        for i in range(len(tokens) - n + 1):
            if any(j in used for j in range(i, i + n)):
                continue
            phrase = " ".join(tokens[i:i + n])
            for alias in by_prefix.get(phrase[:4], []):
                if len(alias.split()) != n:
                    continue
                exact = phrase == alias
                inflected = len(alias) >= 5 and phrase.startswith(alias) and len(phrase) - len(alias) <= 4
                if exact or inflected:
                    found[alias] = aliases[alias]
                    used.update(range(i, i + n))
                    break
    return list(found.values())


def describe_match(match: dict, products: list[dict]) -> dict:
    items = [products[i] for i in match["products"]]
    substances = Counter(tuple(p["subs"]) for p in items if p["subs"]).most_common()
    covered = [p for p in items if p["rep"]]
    example = (covered or items)[0]
    return {
        "label": match["label"],
        "kind": match["kind"],
        "substances": list(substances[0][0]) if len(substances) == 1 else [],
        "ambiguous": len(substances) > 1,
        "substance_variants": [list(s) for s, _ in substances],
        "count": len(items),
        "in_graph": bool(covered),
        "example": example,
    }


BASE_PROMPT = (
    "Odpovedaj po slovensky a iba v rozsahu dostupných dôkazov. "
    "Rozlišuj indikácie, kontraindikácie, upozornenia a interakcie. "
    "Neprítomnosť informácie neznamená neprítomnosť interakcie ani bezpečnosť kombinácie. "
    "Uvádzaj zdroje k tvrdeniam, ak ich kontext poskytuje. Ak dôkaz chýba, priznaj to. "
    "Nehádaj silu, liekovú formu ani zloženie nejednoznačnej značky. "
    "Údaje z reprezentanta rovnakej účinnej látky nemusia pokrývať konkrétny prípravok."
)

# Reserve space for original ADC chunks after graph context and prompt overhead.
QUERY_SETTINGS = {"top_k":40,"chunk_top_k":10,"max_total_tokens":8000,
                  "max_entity_tokens":1000,"max_relation_tokens":1500,"enable_rerank":False}


def prepare_query(question, products, aliases=None, by_prefix=None):
    if aliases is None or by_prefix is None:
        aliases, by_prefix, _ = build_aliases(products)
    matches = [describe_match(m, products) for m in detect_drugs(question, aliases, by_prefix)]
    hints, notes = [], []
    for match in matches:
        if match["ambiguous"]:
            notes.append(f"Názov {match['label']} je nejednoznačný; vyžiadaj presný názov a zloženie.")
        elif match["substances"]:
            hints.append(f"{match['label']} = účinná látka {', '.join(match['substances'])}")
        if not match["in_graph"]:
            ex = match["example"]
            notes.append(f"Pre {ex['name']} nie je v indexe reprezentant s klinickými údajmi; "
                         "príčinou môže byť chýbajúci leták, parsovanie alebo výber dokumentov. "
                         f"Základný údaj ADC: {ex.get('ind','')}. Chýbajúce údaje neznamenajú bezpečnosť.")
    query = question + ("\n\n(Lieky v otázke: " + "; ".join(hints) + ")" if hints else "")
    return {"query":query,"matches":matches,"user_prompt":BASE_PROMPT + "\n" + "\n".join(notes)}