"""
Minimal HTTP client for the BA Agent RL Environment.

Talks to either:
    - a local container (http://localhost:8000)
    - a HuggingFace Space    (https://<user>-<space>.hf.space)

Usage:
    from client import BAAgentClient
    env = BAAgentClient("http://localhost:8000")
    obs = env.reset()
    while not done:
        obs, reward, done, info = env.step(action_type="EXTRACT", payload="...")
"""
from __future__ import annotations

import json
from typing import Any, Dict, List, Optional, Tuple

import requests


class BAAgentClient:
    def __init__(self, base_url: str = "http://localhost:8000", timeout: int = 30):
        self.base_url = base_url.rstrip("/")
        self.timeout = timeout

    def reset(self) -> Dict[str, Any]:
        r = requests.post(f"{self.base_url}/reset", json={}, timeout=self.timeout)
        r.raise_for_status()
        return (r.json() or {}).get("observation", r.json())

    def step(self, action_type: str, payload: str = "") -> Tuple[Dict[str, Any], float, bool, Dict[str, Any]]:
        body = {"action": {"action_type": action_type, "payload": payload}}
        r = requests.post(f"{self.base_url}/step", json=body, timeout=self.timeout)
        r.raise_for_status()
        j = r.json()
        obs = j.get("observation", j)
        reward = float(j.get("reward", 0.0))
        done = bool(j.get("done", False))
        meta = j.get("metadata", {})
        return obs, reward, done, meta

    def state(self) -> Dict[str, Any]:
        r = requests.get(f"{self.base_url}/state", timeout=self.timeout)
        r.raise_for_status()
        return r.json()

    def tasks(self) -> List[Dict[str, Any]]:
        r = requests.get(f"{self.base_url}/api/tasks", timeout=self.timeout)
        r.raise_for_status()
        return r.json()

    def task(self, task_id: str) -> Dict[str, Any]:
        r = requests.get(f"{self.base_url}/api/tasks/{task_id}", timeout=self.timeout)
        r.raise_for_status()
        return r.json()


if __name__ == "__main__":
    env = BAAgentClient()
    obs = env.reset()
    print(f"Task {obs['task_id']} :: {obs['title']}")
    stub_stories = [
        {"title": "Receive document", "description": "As an operator, I want to receive an inbound document, so that I can begin processing.", "acceptance_criteria": "Given inbound, When parsed, Then Document record created."},
        {"title": "Validate document", "description": "As an operator, I want validation, so that bad data is rejected early.", "acceptance_criteria": "Given Document, When validated, Then errors flagged."},
        {"title": "Finalise document", "description": "As an operator, I want to finalise, so that the record is auditable.", "acceptance_criteria": "Given validated Document, When finalised, Then status=Finalised."},
    ]
    plan = [
        ("EXTRACT", "primary workflow, entities=[Actor, System, Document], constraints=compliance."),
        ("INTERVIEW", json.dumps([{"q": "Primary actor?", "a": "Operator"}])),
        ("GRAPH", json.dumps({"nodes": ["Actor", "Document"], "edges": [["Actor", "Document"]]})),
        ("STORY_GEN", json.dumps(stub_stories)),
        ("FINISH", ""),
    ]
    total = 0.0
    for at, pl in plan:
        obs, r, done, meta = env.step(at, pl)
        total += r
        print(f"{at:>10} -> r={r:+.3f} done={done} {meta.get('status','')}")
    print(f"\nEpisode reward = {total:.3f}")
    if "composite_0_to_100" in meta:
        print(f"Composite (bench) = {meta['composite_0_to_100']}/100 ({meta.get('engine')})")