"""A plain tool-calling agent loop over GetFixn, with a hard spending cap. No agent framework.
pip install anthropic requests
export ANTHROPIC_API_KEY=...  ANTHROPIC_MODEL=<a model id you have access to>  GETFIXN_KEY=tfx_live_...
python agent_loop.py "Which of my saved places has the worst county forecast grade today?"
Each grade lookup costs $0.05, each weather lookup $0.0002. The loop stops at MAX_USD (default $1.00).
Any LLM with tool calling works the same way: swap the call in step().
"""
import json
import os
import sys

import gf_client as gf

LOCS = gf.load_locations(os.environ.get("LOCATIONS", "locations.json"))
PRICE = {"forecast_grades": 0.05, "live_weather": 0.0002}  # per location
TOOLS = [
    {"name": "forecast_grades", "description": "county forecast grade A (low) to F (extreme) for saved places. A forecast, not a certainty. $0.05 per place.",
     "input_schema": {"type": "object", "properties": {"places": {"type": "array", "items": {"type": "string"}}}, "required": ["places"]}},
    {"name": "live_weather", "description": "Live temperature and active severe alerts for saved places. $0.0002 per place.",
     "input_schema": {"type": "object", "properties": {"places": {"type": "array", "items": {"type": "string"}}}, "required": ["places"]}},
]


def run_tool(name, args, spent, cap):
    locs = [gf.pick(LOCS, p) for p in args["places"]]
    cost = PRICE[name] * len(locs)
    if spent + cost > cap:
        return "REFUSED: this would exceed the $%.2f budget. Answer with what you already have." % cap, 0.0
    try:
        out = gf.grades(locs) if name == "forecast_grades" else gf.weather(locs)
    except (gf.LowBalance, ValueError, RuntimeError) as e:
        return "ERROR: %s" % e, 0.0
    return json.dumps(out), cost


def step(client, model, messages):
    return client.messages.create(model=model, max_tokens=1024, tools=TOOLS, messages=messages,
                                  system="Storm data assistant. Use the tools. Say it is a forecast and can be wrong. "
                                         "You give data, never trading or investment advice.")


def main(question, client=None, cap=float(os.environ.get("MAX_USD", "1.00"))):
    if client is None:
        import anthropic
        client = anthropic.Anthropic()
    model, spent = os.environ.get("ANTHROPIC_MODEL", ""), 0.0
    messages = [{"role": "user", "content": question}]
    for _ in range(6):  # hard limit on loop turns
        resp = step(client, model, messages)
        if resp.stop_reason != "tool_use":
            return "".join(b.text for b in resp.content if b.type == "text"), spent
        messages.append({"role": "assistant", "content": resp.content})
        results = []
        for b in resp.content:
            if b.type == "tool_use":
                text, cost = run_tool(b.name, b.input, spent, cap)
                spent += cost
                results.append({"type": "tool_result", "tool_use_id": b.id, "content": text})
        messages.append({"role": "user", "content": results})
    return "stopped: too many turns", spent


if __name__ == "__main__":
    answer, spent = main(" ".join(sys.argv[1:]) or "Which saved place has the worst county forecast grade?")
    print(answer, "\n[spent about $%.4f]" % spent)