"""Source-language side: tokenize a controlled English sentence and build the
serial ledger (the decisions reading hides, in reading order).

Nouns come from the lexicon, verbs are relations between meshes, and glue
words (the, a, and) commit no geometry -- determiners open a "which?" question.
A word with no asset is not available unless the caller maps it to a primitive.

Paragraphs are sentences plus bindings. "the <noun>" when that noun is already in
the ledger, and pronouns (it, she, they...), refer back instead of adding an item;
each reference is surfaced as an implied question the discussion can rebind.
"a/an/another <noun>" always adds a new item.
"""
import re

from .lexicon import ASSET_TIERS, LEXICON, PRIMITIVE_SHAPES, TIERS

DETERMINERS = {"the", "a", "an", "another"}
NEW_DETERMINERS = {"a", "an", "another"}  # always introduce a new item
PRONOUNS = {"it", "he", "she", "they", "him", "her", "them"}
# What a pronoun means by default: the most recent object, or the subject of the most
# recent clause (how readers usually take "it": the cat sat on the table; it ... = the cat).
PRONOUN_RULES = ("recent", "subject")
CONJUNCTIONS = {"and", "then"}
COPULAS = {"is", "was", "are", "were"}

# verb -> (relation, implied object when none is stated)
VERBS = {}
for v in ("sat", "sits", "sit", "stood", "stands", "lay", "lies", "rested", "rests"):
    VERBS[v] = ("on", "ground")
for v in ("shone", "shines", "glowed", "glows"):
    VERBS[v] = ("lights", None)  # the subject becomes the light source
for v in ("faced", "faces"):
    VERBS[v] = ("faces", None)

PREPOSITIONS = {
    "in front of": "in_front_of", "next to": "beside", "on top of": "on",
    "on": "on", "under": "under", "beneath": "under", "beside": "beside",
    "near": "near", "behind": "behind", "above": "above",
}

# Nominal heights (m) for assets that carry no real-world scale (Objaverse, primitives).
NOMINAL_HEIGHT_M = {"cat": 0.3, "dog": 0.55, "chair": 0.9, "table": 0.75, "car": 1.5,
                    "tree": 5.0, "person": 1.7, "house": 6.0, "cup": 0.1, "bird": 0.2}
DEFAULT_HEIGHT_M = 0.5


def grammar_words():
    return sorted(DETERMINERS | CONJUNCTIONS | COPULAS | PRONOUNS | set(VERBS) | set(PREPOSITIONS))


def parse_primitives(primitives):
    """{"sun": "sphere"} -> validated dict of word -> shape."""
    out = {}
    for word, shape in (primitives or {}).items():
        if shape not in PRIMITIVE_SHAPES:
            raise ValueError(f"unknown primitive {shape!r} for {word!r}; use one of {PRIMITIVE_SHAPES}")
        out[word.lower().strip()] = shape
    return out


def tokenize(sentence, primitives=None):
    """Greedy longest match over words: prepositions and lexicon phrases can span words."""
    primitives = primitives or {}
    words = re.findall(r"[a-z]+", sentence.lower())
    max_n = max(3, LEXICON.max_phrase_words())
    tokens, i = [], 0
    while i < len(words):
        for n in range(min(max_n, len(words) - i), 0, -1):
            text = " ".join(words[i:i + n])
            kind = classify(text) if n == 1 else classify_phrase(text)
            if kind:
                tokens.append({"index": len(tokens), "text": text, "kind": kind})
                i += n
                break
        else:
            kind = "primitive" if words[i] in primitives else "unknown"
            tokens.append({"index": len(tokens), "text": words[i], "kind": kind})
            i += 1
    return tokens


def classify_phrase(text):
    if text in PREPOSITIONS:
        return "preposition"
    if LEXICON.available(text):
        return "noun"
    return None


def classify(word):
    if word in DETERMINERS:
        return "determiner"
    if word in PRONOUNS:
        return "pronoun"
    if word in CONJUNCTIONS:
        return "conjunction"
    if word in COPULAS:
        return "copula"
    if word in VERBS:
        return "verb"
    if word in PREPOSITIONS:
        return "preposition"
    if LEXICON.available(word):
        return "noun"
    return None


# ---------------- ledger ----------------

POOL = 9  # candidates taken from a tier: the chosen one plus alternatives


def primitive_candidate(shape="box"):
    return {"source": "primitive", "shape": shape}


def _object_decision(word, shape=None):
    """Assets first, in tier order; a primitive only when the caller chose one."""
    if shape:
        return {"tier": "primitive", "chosen": primitive_candidate(shape), "alternatives": [],
                "rejected_tiers": [], "skipped": []}
    found = LEXICON.lookup(word)
    for tier in ASSET_TIERS:
        cands = found.get(tier, [])
        if cands:
            return {"tier": tier, "chosen": cands[0], "alternatives": cands[1:POOL],
                    "rejected_tiers": [], "skipped": []}
    raise ValueError(f"'{word}' has no asset")


def next_tier(decision, word):
    """Rejection: simready -> objaverse -> primitive -> generation."""
    found = LEXICON.lookup(word)
    start = TIERS.index(decision["tier"]) + 1
    for tier in TIERS[start:]:
        if tier in ASSET_TIERS:
            cands = found.get(tier, [])
            if not cands:
                continue
            return {"tier": tier, "chosen": cands[0], "alternatives": cands[1:POOL]}
        if tier == "primitive":
            return {"tier": tier, "chosen": primitive_candidate(), "alternatives": []}
        raise ValueError(f"'{word}' has no tier left after {decision['tier']}: generation is not "
                         f"wired up yet. Use shape or height_m instead, or rephrase.")
    raise ValueError(f"'{word}' has no tier left after {decision['tier']}")


def candidate_label(cand):
    return cand.get("name") or cand.get("category") or cand.get("shape") or cand.get("prompt") or "?"


def _reference_question(qid, token, word, target, others, determiner=None):
    """An implied question for a word that refers back instead of adding an item."""
    if determiner:
        question, alt = f"which {word}? the one already mentioned", [f"say 'another {word}' for a new one"]
    else:
        question, alt = f"'{word}' = the {target['word']}?", []
    return {"id": qid, "token": token["index"], "about": target["id"], "question": question,
            "default_chosen": f"the {target['word']} ({target['id']})", "implied": True,
            "alternatives": alt + [f"the {o['word']} ({o['id']})" for o in others],
            "binds": {"token": token["index"], "to": target["id"]}}


def build_ledger(sentence, primitives=None, prior=None, pronouns="recent"):
    """The ledger for one sentence. With `prior` (an existing ledger), token indices
    continue after it and words can refer back to its objects; the result holds only
    the new tokens, items and questions (extend_ledger merges them). `pronouns` picks
    the default referent (PRONOUN_RULES); the other candidates stay as alternatives."""
    if pronouns not in PRONOUN_RULES:
        raise ValueError(f"pronouns must be one of {list(PRONOUN_RULES)}, got {pronouns!r}")
    primitives = parse_primitives(primitives)
    tokens = tokenize(sentence, primitives)
    unknown = [t["text"] for t in tokens if t["kind"] == "unknown"]
    if unknown:
        raise ValueError(f"Not available (no asset): {', '.join(unknown)}. "
                         f"Rephrase, or map each to a primitive: {PRIMITIVE_SHAPES}")
    offset = len(prior["tokens"]) if prior else 0
    for t in tokens:
        t["index"] += offset
        t["commits"], t["opens"] = [], []

    items, unresolved = [], []
    by_id = {i["id"]: i for i in prior["items"]} if prior else {}
    objects = [i for i in prior["items"] if i["kind"] == "object"] if prior else []  # in reading order
    subjects = [r["subject"] for r in prior["items"] if r["kind"] == "relation"] if prior else []  # clause order

    def add_item(item, token):
        items.append(item)
        by_id[item["id"]] = item
        token["commits"].append(item["id"])
        if item["kind"] == "object":
            objects.append(item)
        return item

    clauses, current = [], []
    for t in tokens:
        if t["kind"] == "conjunction":
            clauses.append(current)
            current = []
        else:
            current.append(t)
    clauses.append(current)

    pending_q = None
    for clause in clauses:
        subject, verb_rel, implied_obj, prep_rel, obj, verb_tok = None, None, None, None, None, None
        subject_tok = obj_tok = None
        for t in clause:
            kind = t["kind"]
            if kind == "determiner":
                pending_q = {"id": f"q{t['index']}", "token": t["index"], "about": None, "determiner": t["text"]}
                t["opens"].append(pending_q["id"])
                unresolved.append(pending_q)
            elif kind in ("noun", "primitive", "pronoun"):
                in_object = subject is not None and (verb_rel is not None or prep_rel is not None)
                if kind == "pronoun":
                    pool = [o for o in objects if not (in_object and o is subject)]
                    if not pool:
                        raise ValueError(f"'{t['text']}' refers to nothing yet: name the object first")
                    item = pool[-1]  # the most recent object it can mean
                    if pronouns == "subject":
                        ids = {o["id"]: o for o in pool}
                        item = next((ids[s] for s in reversed(subjects) if s in ids), item)
                    q = _reference_question(f"q{t['index']}", t, t["text"], item,
                                            [o for o in reversed(pool) if o is not item])
                    q["rule"] = pronouns
                    t["opens"].append(q["id"])
                    unresolved.append(q)
                else:
                    same = [o for o in objects if o["word"] == t["text"]]
                    if same and pending_q is not None and pending_q["determiner"] not in NEW_DETERMINERS:
                        item = same[-1]  # "the cat" again: the cat already mentioned
                        others = [o for o in objects if o is not item][::-1]
                        pending_q.update(_reference_question(pending_q["id"], t, t["text"], item, others,
                                                             determiner=pending_q["determiner"]))
                    else:
                        # height is settled with the asset: metric for simready, nominal otherwise
                        item = add_item({"id": f"n{t['index']}", "kind": "object", "word": t["text"],
                                         "token": t["index"],
                                         "decision": _object_decision(t["text"], primitives.get(t["text"])),
                                         "height_m": {"value": None, "source": "pending"},
                                         "frozen": False}, t)
                        if pending_q is not None:
                            pending_q["about"] = item["id"]
                            pending_q["question"] = f"which {t['text']}?"
                    pending_q = None
                if item["id"] not in t["commits"]:
                    t["refers_to"] = item["id"]
                if in_object:
                    obj, obj_tok = item, t["index"]
                else:
                    subject, subject_tok = item, t["index"]
            elif kind == "verb":
                verb_rel, implied_obj = VERBS[t["text"]]
                verb_tok = t
            elif kind == "copula":
                verb_tok = t
            elif kind == "preposition":
                prep_rel = PREPOSITIONS[t["text"]]
                verb_tok = t
        if subject is None:
            raise ValueError("Each clause needs a subject noun")
        relation = prep_rel or verb_rel
        if relation is None:
            continue
        if obj is None and relation == "on" and implied_obj:
            # Existing needs a reality: the verb implies a support the sentence never
            # names. Visualising exposes it, so it is surfaced as a decision too.
            obj = by_id.get(implied_obj)
            if obj is None:
                obj = add_item({"id": implied_obj, "kind": "environment", "role": "plane", "word": implied_obj,
                                "token": verb_tok["index"], "implied_by": verb_tok["index"],
                                "frozen": False}, verb_tok)
                q = {"id": f"q{verb_tok['index']}i", "token": verb_tok["index"], "about": obj["id"],
                     "question": f"{subject['word']} {verb_tok['text']} on what?",
                     "default_chosen": "ground plane", "implied": True,
                     "alternatives": [f"rephrase: '{verb_tok['text']} on the <noun>'"]}
                verb_tok["opens"].append(q["id"])
                unresolved.append(q)
        if relation == "lights":
            subject["emits_light"] = True
            obj = {"id": "scene"}
        if obj is None:
            raise ValueError(f"'{relation}' needs an object noun")
        if obj is subject:
            raise ValueError(f"'{verb_tok['text']}' would relate the {subject['word']} to itself")
        subjects.append(subject["id"])
        add_item({"id": f"r{verb_tok['index']}", "kind": "relation", "relation": relation,
                  "subject": subject["id"], "object": obj["id"], "token": verb_tok["index"],
                  "tokens": {"subject": subject_tok, "object": obj_tok}, "frozen": False}, verb_tok)

    for q in unresolved:
        q.pop("determiner", None)
        if q["about"] is None:
            q["question"] = "determiner without a noun"
    lit = any(i.get("emits_light") for i in by_id.values())
    asked = bool(prior) and any(q["id"] == "q_light" for q in prior["unresolved"])
    if not lit and not asked:
        # Seeing needs light, as sitting needs a ground: unstated, so surfaced like the ground.
        unresolved.append({"id": "q_light", "token": None, "about": "light", "question": "what light?",
                           "default_chosen": "daylight", "implied": True,
                           "alternatives": ["rephrase: '<noun> shone' or '<noun> glowed'"]})
    return {"version": 1, "sentence": sentence, "primitives": primitives, "pronouns": pronouns,
            "tokens": tokens, "items": items, "unresolved": unresolved}


def extend_ledger(ledger, sentence, primitives=None, pronouns=None):
    """Add a sentence to a ledger: the next sentence of a paragraph. `pronouns`
    defaults to the rule the ledger was built with."""
    prims = {**ledger.get("primitives", {}), **parse_primitives(primitives)}
    pronouns = pronouns or ledger.get("pronouns", "recent")
    new = build_ledger(sentence, prims, prior=ledger, pronouns=pronouns)
    lit = any(i.get("emits_light") for i in ledger["items"] + new["items"])
    unresolved = [q for q in ledger["unresolved"] + new["unresolved"] if not (lit and q["id"] == "q_light")]
    return {**ledger, "sentence": ledger["sentence"].rstrip(" .") + ". " + sentence.strip(),
            "sentences": ledger.get("sentences", [ledger["sentence"]]) + [sentence.strip()],
            "primitives": prims, "pronouns": pronouns, "tokens": ledger["tokens"] + new["tokens"],
            "items": ledger["items"] + new["items"], "unresolved": unresolved}


def check_words(sentence, primitives=None):
    """Per-word status, the words with no asset, and completions for the word being typed."""
    primitives = parse_primitives(primitives)
    tokens = tokenize(sentence, primitives)
    words = []
    for t in tokens:
        entry = {"text": t["text"], "kind": t["kind"]}
        if t["kind"] == "noun":
            entry["available_in"] = LEXICON.counts(t["text"])
            entry["maps_to"] = candidate_label(_object_decision(t["text"])["chosen"])  # what you'd get
        elif t["kind"] == "primitive":
            entry["primitive"] = primitives[t["text"]]
        words.append(entry)
    completions = []
    if sentence and not sentence[-1].isspace():
        last = re.findall(r"[a-z]+", sentence.lower())
        if last:
            prefix = last[-1]
            completions = [w for w in grammar_words() if w.startswith(prefix)][:6]
            completions += LEXICON.complete(prefix, 12)
    missing = [t["text"] for t in tokens if t["kind"] == "unknown"]
    return {"ok": bool(tokens) and not missing, "words": words, "missing": missing,
            "primitive_shapes": PRIMITIVE_SHAPES, "completions": completions}