#!/usr/bin/env python3
"""Extract trigger verbs from the JTBD `situation` field and order them into
universal journey stages.

The "When I ..." clause of a job statement is really the *trigger* — the action
that initiates the job. This script pulls that action out as a verb:

  - Non-LLM part: spaCy finds the leading action verb of each situation
    (lemma, skipping be/have/want/need and auxiliaries), e.g. "Caring for a
    potted plant" -> "care", "Growing succulents" -> "grow".

  - LLM part (one call): the distinct trigger verbs (with counts and a couple of
    example situation->outcome for context) are assigned to a fixed set of
    UNIVERSAL journey stages, so the grouping reflects *where in the customer
    journey* the trigger occurs rather than an invented per-domain category:

        Acquire/Start -> Maintain/Use -> Problem/Breakdown
        -> Diagnose/Resolve -> Optimize/Prevent

    The stages are fixed and universal; only the verb -> stage assignment is
    per-run. Falls back to a small hardcoded verb map if the LLM call fails.

Outputs (alongside the input file):
    <base>__trigger_verbs.csv    comment, trigger_verb
    <base>__trigger_stages.csv   trigger_verb, journey_stage

Usage:
    python extract_trigger_verbs.py <base>__jtbd_archetypes.csv
"""

import argparse
import re
import sys
from collections import Counter
from pathlib import Path

import pandas as pd
import spacy

from llm_client import call_llm_chat_json


JOURNEY_STAGES = [
    "Acquire",
    "Maintain",
    "Problem",
    "Resolve",
    "Improve",
]

SKIP_VERBS = {"be", "have", "do", "want", "need", "will", "would", "can", "could",
              "may", "might", "must", "shall", "should"}

# Fallback verb -> stage map used if the LLM call fails or misses a verb.
# "grow" is ongoing cultivation (not acquisition) — the "when I'm growing X"
# trigger is the Maintain stage, not Acquire.
FALLBACK_MAP = {
    "buy": "Acquire", "purchase": "Acquire", "plant": "Acquire", "sow": "Acquire",
    "start": "Acquire", "build": "Acquire", "prepare": "Acquire", "set": "Acquire",
    "acquire": "Acquire", "adopt": "Acquire", "bring": "Acquire", "install": "Acquire",
    "transplant": "Acquire", "get": "Acquire", "propagate": "Acquire",
    "grow": "Maintain", "care": "Maintain", "manage": "Maintain", "keep": "Maintain",
    "maintain": "Maintain", "water": "Maintain", "tend": "Maintain",
    "use": "Maintain", "live": "Maintain", "house": "Maintain",
    "run": "Maintain", "transfer": "Maintain", "trim": "Maintain", "prune": "Maintain",
    "feed": "Maintain", "monitor": "Maintain",
    "struggle": "Problem", "encounter": "Problem", "deal": "Problem",
    "face": "Problem", "lose": "Problem", "battle": "Problem",
    "suffer": "Problem", "overwater": "Problem", "wilt": "Problem", "appear": "Problem",
    "decide": "Resolve", "diagnose": "Resolve", "check": "Resolve",
    "determine": "Resolve", "assess": "Resolve", "figure": "Resolve",
    "identify": "Resolve", "fix": "Resolve", "solve": "Resolve",
    "remove": "Resolve", "troubleshoot": "Resolve", "repot": "Resolve", "replant": "Resolve",
    "optimize": "Improve", "calibrate": "Improve", "adjust": "Improve",
    "learn": "Improve", "try": "Improve", "experiment": "Improve",
    "improve": "Improve", "measure": "Improve", "track": "Improve",
    "dial": "Improve",
}


def extract_trigger_verb(situation, nlp):
    situation = (situation or "").strip()
    if not situation:
        return ""
    doc = nlp(situation)
    for tok in doc:
        if tok.pos_ in ("VERB", "AUX"):
            lem = tok.lemma_.lower()
            if lem in SKIP_VERBS:
                continue
            return lem
    for tok in doc:
        if tok.pos_ not in ("PUNCT", "SPACE", "DET", "ADP", "PRON", "CCONJ", "SCONJ", "PART") and not tok.is_stop:
            return tok.lemma_.lower()
    return ""


_PHRASE_CUT = re.compile(r',\s*|\s+in\s+|\s+with\s+|\s+while\s+|\s+during\s+|\s+especially\s+|\s+due\s+|\s+on\s+a\s+', re.I)
_LEADING_NOISE = re.compile(r'^(they|the user|user|a person|someone|people|i)\s+', re.I)
_LEADING_COPULA = re.compile(r'^(is|are|was|were|has|have|had|being|having)\s+', re.I)
_TRAILING_STOP = {"the", "a", "an", "to", "of", "for", "in", "on", "at", "and", "or", "with"}


def extract_trigger_phrase(situation):
    """Extract a short trigger phrase (verb + object) from a situation clause,
    e.g. "Growing succulents, especially in a humid climate" -> "growing succulents"."""
    s = (situation or "").strip()
    if not s:
        return ""
    s = _PHRASE_CUT.split(s, maxsplit=1)[0]
    s = _LEADING_NOISE.sub("", s)
    s = _LEADING_COPULA.sub("", s)
    words = s.split()[:5]
    while words and words[-1].lower().rstrip('.,;:') in _TRAILING_STOP:
        words = words[:-1]
    return " ".join(words).strip().lower()


def assign_stages_llm(verbs_info):
    """One LLM call: map each trigger verb to a journey stage. Each verb carries
    its example phrases so the LLM understands the verb's meaning in context
    (e.g. "grow" -> "growing succulents" is Maintain, not Acquire)."""
    verb_block = "\n".join(
        f"{v} ({c}): " + "; ".join(phrases)
        for v, c, phrases in verbs_info
    )

    system_message = (
        "You organize customer research into the stages of a customer journey. "
        "You are given trigger verbs (the actions that initiate a customer's job) "
        "with how often each occurs and the phrases they appear in, so you can see "
        "the context. Assign each verb to the ONE journey stage where that action "
        "typically occurs. Use exactly these stage names: " + ", ".join(JOURNEY_STAGES) + ".\n"
        "Calibration: 'buy' (buying a plant) -> Acquire, 'plant' (planting seeds) -> Acquire, "
        "'grow' (growing succulents) -> Maintain, 'care' (caring for a plant) -> Maintain, "
        "'manage' (managing multiple plants) -> Maintain, 'water' (watering plants) -> Maintain, "
        "'struggle' (struggling with overwatering) -> Problem, 'deal' (dealing with a pest) -> Problem, "
        "'decide' (deciding when to water) -> Resolve, 'diagnose' (diagnosing the problem) -> Resolve, "
        "'check' (checking soil moisture) -> Resolve, "
        "'try' (trying to fix it) -> Improve, 'learn' (learning to water correctly) -> Improve, "
        "'measure' (measuring moisture) -> Improve.\n"
        'Return ONLY this JSON: {"mapping": {"verb": "stage", ...}}'
    )
    user_message = (
        f"Trigger verbs (with example phrases):\n{verb_block}\n\n"
        f"Assign every verb to one of: {', '.join(JOURNEY_STAGES)}"
    )

    try:
        result = None
        for _attempt in range(3):
            result = call_llm_chat_json(
                {"temperature": 0.0, "max_tokens": 1000, "retry_on_invalid_json": 2},
                system_message=system_message,
                user_message=user_message,
            )
            if result and isinstance(result.get("mapping"), dict):
                break
            result = None
    except Exception:
        result = None

    mapping = {}
    if result and isinstance(result.get("mapping"), dict):
        for v, s in result["mapping"].items():
            if isinstance(v, str) and s in JOURNEY_STAGES:
                mapping[v.lower().strip()] = s
    return mapping


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("input", help="Path to <base>__jtbd_archetypes.csv")
    args = parser.parse_args()

    input_path = Path(args.input)
    if not input_path.exists():
        print(f"Error: input file not found: {input_path}")
        sys.exit(1)

    base_name = input_path.name.replace("__jtbd_archetypes.csv", "")
    out_dir = input_path.parent

    df = pd.read_csv(input_path)
    nlp = spacy.load("en_core_web_sm")

    df["trigger_verb"] = df["situation"].fillna("").astype(str).map(lambda s: extract_trigger_verb(s, nlp))
    df["trigger_phrase"] = df["situation"].fillna("").astype(str).map(extract_trigger_phrase)

    verb_counts = Counter(df["trigger_verb"][df["trigger_verb"] != ""])
    # stage the recurring verbs (>=2); singletons are mostly extraction noise
    meaningful_verbs = {v: c for v, c in verb_counts.items() if c >= 2}

    # verb -> example phrases (context for the LLM)
    verb_phrases = {}
    for v in meaningful_verbs:
        ph = df.loc[df["trigger_verb"] == v, "trigger_phrase"].dropna().astype(str)
        verb_phrases[v] = [p for p in dict.fromkeys(ph) if p][:3]

    verbs_info = [(v, c, verb_phrases[v]) for v, c in sorted(meaningful_verbs.items(), key=lambda x: -x[1])]

    llm_mapping = assign_stages_llm(verbs_info)
    if llm_mapping:
        print(f"  LLM mapped {len(llm_mapping)} verbs to journey stages.")
    else:
        print("  LLM call failed; using fallback verb map.")

    stage_by_verb = {}
    for v in verb_counts:
        stage_by_verb[v] = llm_mapping.get(v, FALLBACK_MAP.get(v, ""))

    # phrase -> stage, derived from the phrase's verb
    phrase_to_verb = (df.drop_duplicates("trigger_phrase")
                        .set_index("trigger_phrase")["trigger_verb"].to_dict())
    stage_by_phrase = {p: stage_by_verb.get(v, "") for p, v in phrase_to_verb.items()}

    df[["comment", "trigger_verb", "trigger_phrase"]].to_csv(out_dir / f"{base_name}__trigger_verbs.csv", index=False)

    stage_rows = [{"trigger_phrase": p, "journey_stage": s} for p, s in stage_by_phrase.items()]
    pd.DataFrame(stage_rows).to_csv(out_dir / f"{base_name}__trigger_stages.csv", index=False)

    print(f"  Saved -> {base_name}__trigger_verbs.csv")
    print(f"  Saved -> {base_name}__trigger_stages.csv")
    print("\nTrigger phrases -> journey stage:")
    by_stage = {}
    for p, s in stage_by_phrase.items():
        by_stage.setdefault(s or "(unassigned)", []).append(p)
    for s in JOURNEY_STAGES + ["(unassigned)"]:
        if s in by_stage:
            # show the most common phrases only
            shown = sorted(by_stage[s], key=lambda x: -len(df[df["trigger_phrase"] == x]))[:12]
            print(f"  {s}: {', '.join(shown)}")

    print("\nDone.")


if __name__ == "__main__":
    main()