import json
import gradio as gr
from engine import walk, rationale
import ai_layer as A

INTRO = ("Describe your footwear product in your own words. "
         "Include the materials, the type, who it is for, and the approximate value per pair. "
         "I read your description, then ask only for what the schedule still needs.")


def fmt_result(res, facts):
    lines = []
    lines.append(f"Suggested HTS code {res['code']}")
    lines.append(f"General duty rate {res['general'] or 'see schedule'}")
    lines.append("")
    lines.append("Why this code")
    for i, step in enumerate(res["path"], 1):
        num = step["htsno"] or "grouping"
        lines.append(f"  {i}. {num}  {step['desc']}")
    lines.append("")
    lines.append("Basis. HTS Chapter 64, 2026 revision, walked per the chapter notes.")
    lines.append("Decision support only. The importer or a licensed broker makes the final decision.")
    lines.append("")
    lines.append("Describe another product to classify it.")
    return "\n".join(lines)


def fmt_review(reason, at=None):
    lines = ["This case is routed to expert review."]
    lines.append(f"Reason. {reason}.")
    if at:
        lines.append(f"The unresolved condition. {at}")
    lines.append("A licensed broker should make this call. "
                 "Describe another product to classify it.")
    return "\n".join(lines)


def advance(client, state):
    """Walk with current facts. Returns the assistant message and whether done."""
    r = walk(state["facts"])
    if r["status"] == "code":
        msg = fmt_result(r, state["facts"])
        return msg, True
    if r["status"] == "need_fact":
        try:
            q = A.ask_gap(client, r["fact"])
        except Exception:
            q = f"Please tell me {A.FACT_KEYS[r['fact']]}."
        state["pending"] = ("fact", r["fact"])
        return q, False
    if r["status"] == "review" and r.get("at"):
        try:
            q = A.ask_branch(client, r["at"])
        except Exception:
            q = f"Does this condition apply to your product. {r['at']}"
        state["pending"] = ("branch", r["at"])
        return q, False
    return fmt_review(r.get("reason", "unresolved")), True


def respond(message, history, state):
    state = dict(state or {"facts": {}, "pending": None, "retried": False})
    history = history + [{"role": "user", "content": message}]

    try:
        client = A.get_client()
    except Exception:
        history.append({"role": "assistant",
                        "content": "The model is not configured. The Space owner must set the GROQ_API_KEY secret."})
        return history, state, ""

    try:
        if state["pending"] is None:
            state["facts"] = A.extract_facts(client, message)
            state["facts"]["_desc"] = message
            msg, done = advance(client, state)
        else:
            kind, key = state["pending"]
            if kind == "fact":
                got = A.read_answer(client, key, message)
                if got:
                    state["facts"].update(got)
                    state["pending"], state["retried"] = None, False
                    msg, done = advance(client, state)
                elif not state["retried"]:
                    state["retried"] = True
                    msg, done = "Please answer that directly so the classification stays accurate.", False
                else:
                    msg, done = fmt_review(f"the fact could not be established, {A.FACT_KEYS[key]}"), True
            else:
                holds = A.read_branch_answer(client, key, message)
                if holds is None:
                    msg, done = fmt_review("a schedule condition needs expert judgment", key), True
                else:
                    state["facts"].setdefault("_overrides", {})[key] = holds
                    state["pending"] = None
                    msg, done = advance(client, state)
        if done:
            state = {"facts": {}, "pending": None, "retried": False}
    except Exception as e:
        msg = f"The model call failed. Reason {e}. Send your message again in a moment."

    history.append({"role": "assistant", "content": msg})
    return history, state, ""


with gr.Blocks(title="TariffWise") as demo:
    gr.Markdown("# TariffWise")
    gr.Markdown("AI tariff classification for footwear. A model reads your description and runs the interview. "
                "A deterministic engine walks the real 2026 HTS Chapter 64 and assigns the code. "
                "Every result carries its full derivation.")
    chatbot = gr.Chatbot(height=460, value=[{"role": "assistant", "content": INTRO}])
    st = gr.State(None)
    box = gr.Textbox(placeholder="Describe your product, or answer the question", show_label=False)
    box.submit(respond, [box, chatbot, st], [chatbot, st, box])
    gr.Markdown("Decision support only. Not legal advice.")

if __name__ == "__main__":
    demo.launch()