import json, re
import gradio as gr

try:
    from transformers import pipeline
except ImportError:  # pragma: no cover - handled gracefully if transformers missing
    pipeline = None

# Load knowledge
with open("kb.json", "r", encoding="utf-8") as f:
    KB = json.load(f)

def find_cards(topic, home_type=None):
    cards = [c for c in KB if c["topic"]==topic]
    if home_type:
        cards = [c for c in cards if home_type in c["context"]]
    return cards

def short_answer(card):
    fact = card["fact"]
    tip = card["actions"][0]
    return f"{card['title']}\n{fact}\nAction: {tip}"

GUIDE_QS = [
    ("city", "Pick your city", ["Karachi","Lahore","Islamabad","Other"]),
    ("home", "Where do you live?", ["Apartment","House","Hostel/School"]),
    ("focus", "What do you want to improve?", ["Water","Electricity","Waste"]),
    ("habit", "Pick one habit to fix", {
        "Water":["Long showers","Leaky tap","Dishwashing"],
        "Electricity":["AC at night","Old bulbs","Many chargers"],
        "Waste":["No sorting","E-waste lying around","Food scraps"]
    }),
    ("extras", "Extras available? (pick any, or 'None')", ["Fan","Washing machine","Garden","Rooftop","None"])
]

INTENT_LABELS = [
    "water",
    "electricity",
    "waste",
    "greeting",
    "smalltalk",
    "affirmation",
    "gratitude",
    "checklist",
    "goodbye"
]
INTENT_THRESHOLD = 0.45
_INTENT_PIPELINE = None
_INTENT_PIPELINE_ERR = None


def load_intent_classifier():
    """Lazily load the zero-shot intent classifier pipeline."""
    global _INTENT_PIPELINE, _INTENT_PIPELINE_ERR
    if pipeline is None:
        _INTENT_PIPELINE_ERR = "transformers not installed"
        return None
    if _INTENT_PIPELINE is not None or _INTENT_PIPELINE_ERR:
        return _INTENT_PIPELINE
    try:
        _INTENT_PIPELINE = pipeline(
            "zero-shot-classification",
            model="typeform/distilbert-base-uncased-mnli"
        )
    except Exception as err:  # pragma: no cover - model download/setup issues
        _INTENT_PIPELINE_ERR = err
        return None
    return _INTENT_PIPELINE


def classify_intent(text):
    """Return (label, score) for the strongest intent, or (None, 0.0)."""
    if not text:
        return None, 0.0
    clf = load_intent_classifier()
    if clf is None:
        return None, 0.0
    result = clf(text, candidate_labels=INTENT_LABELS, multi_label=True)
    labels = result.get("labels") or []
    scores = result.get("scores") or []
    if not labels or not scores:
        return None, 0.0
    best_label, best_score = max(zip(labels, scores), key=lambda pair: pair[1])
    if best_score < INTENT_THRESHOLD:
        return None, best_score
    return best_label, best_score


def map_habit_to_cards(focus, habit):
    focus = focus.lower()
    if focus=="water":
        if "leak" in habit.lower(): return ["water_leak","water_bottle"]
        if "shower" in habit.lower(): return ["water_shower","water_bottle"]
        return ["water_shower","water_leak"]
    if focus=="electricity":
        if "ac" in habit.lower(): return ["power_ac26","power_fanfirst","power_led"]
        if "bulb" in habit.lower(): return ["power_led","power_ac26"]
        return ["power_fanfirst","power_led"]
    if focus=="waste":
        if "e-waste" in habit.lower(): return ["waste_ewaste","waste_bins"]
        if "sorting" in habit.lower(): return ["waste_bins","waste_rinse"]
        return ["waste_bins","waste_rinse"]
    return []

def estimate_impact(cards):
    L = sum(c.get("impact",{}).get("daily_L",0) for c in cards)
    K = sum(c.get("impact",{}).get("kwh_day",0) for c in cards)
    W = sum(c.get("impact",{}).get("kg_week",0) for c in cards)
    parts=[]
    if L: parts.append(f"water saved ~{int(L)} L/day")
    if K: parts.append(f"power cut ~{K:.1f} kWh/day")
    if W: parts.append(f"waste diverted ~{W:.1f} kg/week")
    return "; ".join(parts) if parts else "visible within a week"

def init_state():
    return {"mode":"tutor","g_step":0,"g_answers":{}}

def bot_reply(user, state):
    # Mode switches
    txt = (user or "").strip().lower()
    if txt in ["tutor","/tutor"]:
        state["mode"]="tutor"
        return "Tutor on. Ask me anything about water, electricity, or waste.", state
    if txt in ["guide","/guide","plan"]:
        state["mode"]="guide"; state["g_step"]=0; state["g_answers"]={}
        key, q, opts = GUIDE_QS[0]
        return f"{q}: {', '.join(opts)}", state

    # GUIDE mode flow
    if state["mode"]=="guide":
        step = state["g_step"]
        key, q, opts = GUIDE_QS[step]
        ans = (user or "").strip()

        if isinstance(opts, list):
            if not any(ans.lower()==o.lower() for o in opts):
                return f"Please choose one of: {', '.join(opts)}", state
            state["g_answers"][key]=ans
        else:
            focus = state["g_answers"].get("focus","Water")
            options = opts[focus]
            if not any(ans.lower()==o.lower() for o in options):
                return f"Please choose one of: {', '.join(options)}", state
            state["g_answers"][key]=ans

        state["g_step"] += 1
        if state["g_step"] >= len(GUIDE_QS):
            city = state["g_answers"]["city"]
            home = state["g_answers"]["home"].lower().replace("hostel/school","school")
            focus = state["g_answers"]["focus"]
            habit = state["g_answers"]["habit"]
            card_ids = map_habit_to_cards(focus, habit)
            cards = [c for c in KB if c["id"] in card_ids and (home in c["context"])]
            if not cards:
                cards = find_cards(focus.lower())
            actions = [f"- {c['title']}: {c['actions'][0]}" for c in cards[:3]]
            impact = estimate_impact(cards[:3])
            msg = (
                f"You said: {city}, {home}, focus = {focus}.\n"
                f"Your 3 actions:\n" + "\n".join(actions) +
                f"\nEstimated impact: {impact}\n"
                "Want a 7-day checklist? (type: water / power / waste)"
            )
            state["mode"]="tutor"
            return msg, state
        else:
            key, q, opts = GUIDE_QS[state["g_step"]]
            if isinstance(opts, list):
                return f"{q}: {', '.join(opts)}", state
            else:
                focus = state["g_answers"].get("focus","Water")
                return f"{q}: {', '.join(opts[focus])}", state

    # TUTOR mode quick topic guess
    t = (user or "").lower()
    normalized = re.sub(r"[^\w\s]", "", t).strip()
    intent_label, intent_score = classify_intent(user or "")
    if any(w in t for w in ["water","shower","tap","leak","plant","tree","bottle"]):
        topic="water"
    elif any(w in t for w in ["ac","power","electric","fan","bulb","led"]):
        topic="electricity"
    elif any(w in t for w in ["waste","trash","recycle","bin","e-waste","ewaste","plastic"]):
        topic="waste"
    else:
        topic=None

    if topic is None and intent_label in {"water","electricity","waste"}:
        topic = intent_label

    if topic:
        home_hint = "apartment" if "apartment" in t else ("house" if "house" in t else None)
        cards = find_cards(topic, home_hint) or find_cards(topic)
        card = cards[0]
        return short_answer(card), state

    # small talk and acknowledgements
    if intent_label == "greeting" and intent_score >= INTENT_THRESHOLD:
        return "Salaam! I'm your Green Guide bot - ask me about water, electricity, or waste habits anytime.", state
    if re.search(r"\b(hi|hello|hey|salaam|assalam|assalamu alaikum|assalamualaikum)\b", t):
        return "Salaam! I'm your Green Guide bot - ask me about water, electricity, or waste habits anytime.", state
    if intent_label == "smalltalk" and intent_score >= INTENT_THRESHOLD:
        return "Happy to chat! Got any water, electricity, or waste question for me?", state
    if "how are you" in t:
        return "I'm running on code and clean energy facts! How can I help with your sustainability goal?", state
    if intent_label in {"affirmation","gratitude"} and intent_score >= INTENT_THRESHOLD:
        return "Noted! Tell me what you want to improve: water, electricity, or waste, and I'll share tips.", state
    if normalized in {"ok","okay","yes","yup","sure","cool","great","thanks","thank you","thankyou","nice"}:
        return "Noted! Tell me what you want to improve: water, electricity, or waste, and I'll share tips.", state
    if intent_label == "goodbye" and intent_score >= INTENT_THRESHOLD:
        return "Khuda hafiz! Come back anytime you need another green tip.", state

    # checklists
    if intent_label == "checklist" and intent_score >= INTENT_THRESHOLD:
        return "Need a checklist? Type **water**, **power**, or **waste** and I'll drop a 7-day plan.", state
    if re.search(r"\b(checklist|pledge)\b", t):
        if "power" in t:
            return "Power Ninja 7-day: LED swap ? 26?C nights ? fan-first ? unplug at 10pm ? kill standby ? use daylight ? share result", state
        if "waste" in t:
            return "Waste Warrior 7-day: 2-bin ? rinse items ? paper stack ? e-waste box ? no plastic bag ? find drop-off ? share photo", state
        return "Water Saver 7-day: fix drip · 5-min showers · reuse rinse for plants · full loads · cold bottle in fridge · log habits · share one tip", state

    return "Try **Tutor**: ask me anything about water/electricity/waste. Or type **Guide** for a mini plan.", state

def respond(message, history, state):
    reply, state = bot_reply(message, state)
    history = history + [(message, reply)]
    return history, state

with gr.Blocks(fill_height=True, css="footer{visibility:hidden}") as demo:
    gr.Markdown("### **Green Guide** — Tutor & Guide for a cleaner Pakistan.\nType **Guide** to get a mini plan.")
    chatbot = gr.Chatbot(height=420)
    txt = gr.Textbox(placeholder="Type here… (try: Guide)", autofocus=True)
    state = gr.State(init_state())

    # greet
    chatbot.value = [("", "Salaam! Pick a mode: **Tutor** (ask anything) or **Guide** (I’ll make you a mini plan).")]

    txt.submit(respond, [txt, chatbot, state], [chatbot, state]).then(
        lambda: "", None, [txt]
    )

demo.launch()