import sys, json, re, argparse, gradio as gr
from crop_tools import execute
ap = argparse.ArgumentParser(); ap.add_argument("--backend", default="ollama", choices=["ollama","hf"])
ap.add_argument("--model", default="commodities-crop-analyst"); ap.add_argument("--adapter", default="."); ap.add_argument("--base", default="Qwen/Qwen2.5-7B-Instruct")
args = ap.parse_args()
import datetime
SYSTEM = (open("system_prompt.txt").read() if args.backend == "ollama" else [json.loads(l) for l in open("data/val_v2.jsonl")][0]["messages"][0]["content"]) + f"\nToday's date is {datetime.date.today().isoformat()}. Resolve relative dates (last year, a year ago, last month) from it. Price data runs 2017-01-03 to {open('data/futures_daily.parquet','rb') and __import__('pandas').read_parquet('data/futures_daily.parquet').date.max()}."

if args.backend == "hf":
    import torch
    from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
    from peft import PeftModel
    tok = AutoTokenizer.from_pretrained(args.base)
    model = PeftModel.from_pretrained(AutoModelForCausalLM.from_pretrained(args.base, quantization_config=BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16), device_map="auto"), args.adapter).eval()
    def gen(msgs, n=300):
        ids = tok.apply_chat_template(msgs, add_generation_prompt=True, return_tensors="pt", return_dict=True)["input_ids"].to(model.device)
        return tok.decode(model.generate(ids, max_new_tokens=n, do_sample=False)[0][ids.shape[1]:], skip_special_tokens=True).strip()
else:
    import requests
    def gen(msgs, n=300):
        r = requests.post("http://localhost:11434/api/chat", json={"model": args.model, "messages": msgs, "stream": False, "options": {"temperature": 0, "num_predict": n}}, timeout=300)
        r.raise_for_status(); return r.json()["message"]["content"].strip()

def answer(user, history):
    msgs = [{"role":"system","content":SYSTEM}]
    for h in history[-6:]:
        msgs.append({"role":h["role"],"content":h["content"] if isinstance(h["content"], str) else "[image]"})
    msgs.append({"role":"user","content":user})
    first = gen(msgs, 160)
    m = re.search(r"<tool_call>(.*?)</tool_call>", first, re.S)
    if not m: return first, None, None
    try: call = json.loads(m.group(1))
    except Exception: return "I produced a malformed tool call; please rephrase.", None, first
    res = execute(call)
    msgs += [{"role":"assistant","content":first}, {"role":"tool","content":json.dumps(res)}]
    final = gen(msgs, 300)
    img = res.get("image") if isinstance(res, dict) else None
    return final, img, json.dumps({"call": call, "result": res}, indent=1, default=str)

def chat(user, history):
    text, img, trace = answer(user, history)
    out = [{"role":"assistant","content":text}]
    if img: out.append({"role":"assistant","content":{"path": img}})
    if trace: out.append({"role":"assistant","content":"```json\n" + trace[:1500] + "\n```", "metadata":{"title":"tool call"}})
    return out

gr.ChatInterface(chat, title="Commodities crop analyst (edoggy)",
    description="Answers from USDA CDL area, census, IBGE yields and futures data via tools. Forecasts are statistical baselines with error bars. Not advice.",
    examples=["How much corn did Iowa have in 2023?", "Which Louisiana parish has the most sugarcane?", "Forecast sugar 60 days out",
              "Show me the sugarcane map for Iberia, LA", "Chart orange juice since 2022", "Is Louisiana cane expanding next year?"]
).launch(server_name="0.0.0.0", server_port=7860)