"""
Lab 08c: Give a model two tools, then watch it (ReAct loop)
Chapter 8 - Applied Cybersecurity: Foundations to Mastery

Requires Ollama running on the host at http://localhost:11434.
See Lab 08c (Ollama setup) in the README before running this.

AUTHORISATION GUARD RAIL (mandatory per book):
  The 'fetch_url' tool is restricted to an allow-listed set of local URLs only.
  Do NOT expand the allowlist to arbitrary internet addresses.
  Only test against the provided lab target or a system you own.
"""

import json
import math
import requests

# ── Configuration ─────────────────────────────────────────────────────────────
OLLAMA_BASE   = "http://localhost:11434"
MODEL         = "llama3.2"   # ollama pull llama3.2 first

# AUTHORISATION: only these local addresses are permitted for the fetch_url tool
URL_ALLOWLIST = {
    "http://localhost:8080/lab-data.json",
    "http://127.0.0.1:8080/lab-data.json",
}

# ── Tool definitions ──────────────────────────────────────────────────────────
TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "calculator",
            "description": "Evaluate a simple arithmetic expression and return the result.",
            "parameters": {
                "type": "object",
                "properties": {
                    "expression": {
                        "type": "string",
                        "description": "A Python-safe arithmetic expression, e.g. '42 * 1.15'"
                    }
                },
                "required": ["expression"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "fetch_url",
            "description": (
                "Fetch the content of a URL and return it as text. "
                "RESTRICTED: only allow-listed local URLs are permitted."
            ),
            "parameters": {
                "type": "object",
                "properties": {
                    "url": {
                        "type": "string",
                        "description": "The URL to fetch (must be in the allow-list)."
                    }
                },
                "required": ["url"]
            }
        }
    }
]


# ── Tool implementations ──────────────────────────────────────────────────────
def run_calculator(expression: str) -> str:
    """Evaluate a restricted arithmetic expression."""
    # Restrict to safe characters only - no imports, no builtins
    allowed = set("0123456789.+-*/(). ")
    if not all(c in allowed for c in expression):
        return f"ERROR: expression contains disallowed characters: {expression!r}"
    try:
        result = eval(expression, {"__builtins__": {}}, {"sqrt": math.sqrt, "pow": pow})
        return str(result)
    except Exception as e:
        return f"ERROR: {e}"


def run_fetch_url(url: str) -> str:
    """Fetch a URL - only allow-listed local addresses permitted."""
    if url not in URL_ALLOWLIST:
        return (
            f"BLOCKED: URL '{url}' is not in the allow-list. "
            f"Permitted: {sorted(URL_ALLOWLIST)}"
        )
    try:
        resp = requests.get(url, timeout=3)
        return resp.text
    except requests.exceptions.ConnectionError:
        # Lab data server not running - return synthetic data so the demo works
        print(f"  [tool] lab data server not running at {url}, using synthetic data")
        return json.dumps({"port_scan_hits": 47, "scan_duration_seconds": 30})
    except Exception as e:
        return f"ERROR: {e}"


TOOL_DISPATCH = {
    "calculator": run_calculator,
    "fetch_url":  run_fetch_url,
}


# ── ReAct loop ────────────────────────────────────────────────────────────────
def run_agent(user_prompt: str, max_turns: int = 6):
    """Run the model in a tool-calling loop and print each step."""
    messages = [{"role": "user", "content": user_prompt}]

    print(f"\n{'='*60}")
    print(f"PROMPT: {user_prompt}")
    print(f"{'='*60}")

    for turn in range(max_turns):
        resp = requests.post(
            f"{OLLAMA_BASE}/api/chat",
            json={
                "model":    MODEL,
                "messages": messages,
                "tools":    TOOLS,
                "stream":   False,
            },
            timeout=60,
        )
        resp.raise_for_status()
        data       = resp.json()
        msg        = data["message"]
        tool_calls = msg.get("tool_calls") or []

        if not tool_calls:
            print(f"\n[Turn {turn+1}] FINAL ANSWER:")
            print(msg.get("content", "(no content)"))
            break

        # Print and execute each tool call (the "action" in ReAct)
        messages.append(msg)
        for tc in tool_calls:
            fn_name = tc["function"]["name"]
            fn_args = tc["function"]["arguments"]
            if isinstance(fn_args, str):
                fn_args = json.loads(fn_args)

            print(f"\n[Turn {turn+1}] TOOL CALL: {fn_name}({fn_args})")
            result = TOOL_DISPATCH[fn_name](**fn_args)
            print(f"[Turn {turn+1}] OBSERVATION: {result}")

            messages.append({
                "role":    "tool",
                "content": result,
            })
    else:
        print("\n[max turns reached]")


# ── Main ──────────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    # Check Ollama is reachable first
    try:
        requests.get(f"{OLLAMA_BASE}/api/tags", timeout=3).raise_for_status()
    except Exception:
        print(f"ERROR: Ollama not reachable at {OLLAMA_BASE}")
        print("Run 'ollama serve' and 'ollama pull llama3.2' first (Lab 08c setup).")
        raise SystemExit(1)

    # The book's example: fetch a number from the lab page, then compute from it
    prompt = (
        "Fetch the port scan data from http://localhost:8080/lab-data.json, "
        "then calculate the scan hit rate as hits per second (round to 2 decimal places). "
        "Show your working."
    )
    run_agent(prompt)
