#!/usr/bin/env python3
"""
ONI Agent API Client
====================
CLI tool for interacting with the ONI Agent Bridge Mod (v2).
All commands return structured JSON data.

Usage:
  python oni_api.py <command> [args...]

Commands:
  # Status
  health             Check mod connection
  status             Full game overview
  resources          All resource amounts
  buildings          All buildings
  duplicants         Duplicant details
  research           Tech progress
  rooms              Room layout

  # Map data
  cell <x> <y>       Single cell details
  cells <x> <y> <w> <h>  Rectangular area
  slice <axis> <index> <start> <end>  Row/column scan
  gas <x> <y> <r>    Gas distribution

  # Registry
  registry buildings [filter]  Building definitions
  registry elements [filter]   Element definitions
  registry techs               Tech tree

  # Events
  events [since]     Poll events since sequence number

  # Actions
  dig <x> <y> <w> <h>         Dig area
  build <buildingId> <x> <y>   Build structure
  deconstruct <x> <y>          Deconstruct
  prioritize <x> <y> <p>       Set priority (1-9)
  research <techId>            Set active research
  mop <x> <y>                  Mop liquid
  harvest <x> <y>              Harvest plant
  pause [reason]               Pause game
  unpause [speed]              Resume game
  speed <1|2|3>                Set game speed
  batch <json_file>            Execute batch operations
  camera <x> <y> [zoom]        Move camera
  save [name]                  Save game
  load <name>                  Load save
  priority_global <target> <p> Set global priority
  priority_type <type> <p>     Set type priority

  # Utility
  explore <x> <y> <w> <h>     AI-friendly area summary
"""

import json, os, sys, urllib.request, urllib.error

# Fix GBK encoding on Windows CJK consoles
if sys.stdout.encoding and sys.stdout.encoding.upper() in ('GBK', 'GB2312', 'CP936'):
    import io
    sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8', errors='replace')

CONFIG_PATH = os.path.join(os.path.dirname(os.path.dirname(__file__)), 'config.json')

def load_config():
    with open(CONFIG_PATH) as f:
        return json.load(f)

def api_get(endpoint):
    cfg = load_config()
    url = f"http://{cfg['modHost']}:{cfg['modPort']}{endpoint}"
    try:
        req = urllib.request.Request(url, method='GET')
        with urllib.request.urlopen(req, timeout=cfg['timeout']) as resp:
            return json.loads(resp.read().decode())
    except urllib.error.URLError as e:
        return {"success": False, "error": "connection_error", "errorMessage": str(e)}

def api_post(endpoint, data):
    cfg = load_config()
    url = f"http://{cfg['modHost']}:{cfg['modPort']}{endpoint}"
    try:
        req = urllib.request.Request(
            url, data=json.dumps(data).encode(),
            headers={'Content-Type': 'application/json'},
            method='POST'
        )
        with urllib.request.urlopen(req, timeout=cfg['timeout']) as resp:
            return json.loads(resp.read().decode())
    except urllib.error.URLError as e:
        return {"success": False, "error": "connection_error", "errorMessage": str(e)}

def get_data(endpoint):
    """Get the data field from a response, handling the wrapper."""
    r = api_get(endpoint)
    if r.get("success"):
        return r.get("data", r)
    print(f"Error: {r.get('errorMessage', r.get('error', 'unknown'))}", file=sys.stderr)
    return None

def post_action(endpoint, data=None):
    """Post and print result."""
    r = api_post(endpoint, data or {})
    if r.get("success"):
        d = r.get("data", {})
        if d:
            print(json.dumps(d, indent=2, ensure_ascii=False))
        else:
            print("OK")
    else:
        print(f"Error: {r.get('errorMessage', r.get('error', 'unknown'))}", file=sys.stderr)
    return r

# ── Commands ─────────────────────────────────────────────────

def cmd_health(args):
    r = api_get('/health')
    # Support both new format {success, data} and old format {status}
    if r.get("success"):
        d = r.get("data", r)
        print(f"Status: {d.get('status', '?')}")
        print(f"Version: {d.get('version', 'unknown — old mod, deploy v2')}")
    elif r.get("status") == "ok":
        print(f"Status: ok (old mod v1 — restart game to load v2)")
    else:
        print(f"Connection failed: {r}")

def cmd_status(args):
    game = get_data('/api/state/game')
    if not game: return
    resources = get_data('/api/state/resources') or []
    dups = get_data('/api/state/duplicants') or []
    alerts = get_data('/api/state/alert') or []

    print(f"=== Game ===")
    print(f"  Cycle:      {game.get('cycle', '?')}")
    print(f"  Dupes:      {game.get('duplicantCount', '?')}")
    print(f"  Grid:       {game.get('gridWidth', '?')} x {game.get('gridHeight', '?')}")
    print(f"  Paused:     {game.get('isPaused', '?')}")
    print(f"  Speed:      {game.get('gameSpeed', '?')}x")

    print(f"\n=== Top Resources ===")
    if isinstance(resources, list):
        for r in sorted(resources, key=lambda x: x.get('amountKg', 0), reverse=True)[:20]:
            icon = {'solid': '■', 'liquid': '≈', 'gas': '◌'}.get(r.get('state', ''), '?')
            print(f"  {icon} {r.get('name', '?'):20s} {r.get('amountKg', 0):>10.1f} kg")

    print(f"\n=== Duplicants ===")
    for d in dups if isinstance(dups, list) else []:
        print(f"  {d.get('name', '?'):12s} at ({d.get('x', '?'):4d},{d.get('y', '?'):4d})  hp={d.get('health', 0)}")

    if alerts:
        print(f"\n=== Alerts ===")
        for a in alerts:
            print(f"  [{a.get('severity', '?')}] {a.get('title', '?')}")

def cmd_resources(args):
    data = get_data('/api/state/resources')
    if data:
        print(json.dumps(data, indent=2, ensure_ascii=False))

def cmd_buildings(args):
    data = get_data('/api/state/buildings')
    if data:
        print(json.dumps(data, indent=2, ensure_ascii=False))

def cmd_duplicants(args):
    data = get_data('/api/state/duplicants')
    if data:
        print(json.dumps(data, indent=2, ensure_ascii=False))

def cmd_research(args):
    data = get_data('/api/state/research')
    if data:
        print(json.dumps(data, indent=2, ensure_ascii=False))

def cmd_rooms(args):
    data = get_data('/api/state/rooms')
    if data:
        print(json.dumps(data, indent=2, ensure_ascii=False))

def cmd_cell(args):
    if len(args) < 2: print("Usage: cell <x> <y>", file=sys.stderr); return
    x, y = int(args[0]), int(args[1])
    r = api_get(f'/api/state/cell?x={x}&y={y}')
    if r.get("success"):
        d = r.get("data", {})
        print(f"Cell ({d.get('x')}, {d.get('y')}):")
        print(f"  Element:    {d.get('element', '?')} ({d.get('elementId', '?')})")
        print(f"  Mass:       {d.get('massKg', 0):.2f} kg")
        print(f"  Temp:       {d.get('temperatureC', 0):.1f} °C")
        print(f"  State:      {'Solid' if d.get('isSolid') else 'Liquid' if d.get('isLiquid') else 'Gas' if d.get('isGas') else 'Vacuum'}")
        print(f"  Building:   {d.get('buildingName', 'None')}")
        print(f"  Diggable:   {d.get('isDiggable', False)}")
        print(f"  Duplicant:  {d.get('hasDuplicant', False)}")
    else:
        print(f"Error: {r.get('errorMessage', r.get('error', 'unknown'))}", file=sys.stderr)

def cmd_cells(args):
    if len(args) < 4: print("Usage: cells <x> <y> <width> <height>", file=sys.stderr); return
    x, y, w, h = int(args[0]), int(args[1]), int(args[2]), int(args[3])
    r = api_get(f'/api/state/cells?x={x}&y={y}&width={w}&height={h}')
    if r.get("success"):
        d = r.get("data", {})
        cells = d.get("cells", [])
        print(f"{len(cells)} cells in ({x},{y})-({x+w},{y+h}):")
        for c in cells:
            icon = '■' if c.get('isSolid') else '≈' if c.get('isLiquid') else '◌' if c.get('isGas') else ' '
            b = '⚙' if c.get('hasBuilding') else ' '
            d = '🧑' if c.get('hasDuplicant') else ' '
            print(f"  ({c['x']:3d},{c['y']:3d}) {icon} {c.get('element','?'):15s} {c.get('massKg',0):6.1f}kg {c.get('temperatureC',0):5.1f}°C{b}{d}")
    else:
        print(f"Error: {r.get('errorMessage', r.get('error', 'unknown'))}", file=sys.stderr)

def cmd_slice(args):
    if len(args) < 3: print("Usage: slice <x|y> <index> <start> <end>", file=sys.stderr); return
    axis, idx, start, end = args[0], int(args[1]), int(args[2]), int(args[3])
    r = api_get(f'/api/state/cells/slice?axis={axis}&index={idx}&start={start}&end={end}')
    if r.get("success"):
        d = r.get("data", {})
        cells = d.get("cells", [])
        for c in cells:
            icon = '■' if c.get('isSolid') else '≈' if c.get('isLiquid') else '◌' if c.get('isGas') else ' '
            print(f"  ({c['x']:3d},{c['y']:3d}) {icon} {c.get('element','?'):15s} {c.get('massKg',0):6.1f}kg {c.get('temperatureC',0):5.1f}°C")
    else:
        print(f"Error: {r.get('errorMessage', r.get('error', 'unknown'))}", file=sys.stderr)

def cmd_gas(args):
    if len(args) < 3: print("Usage: gas <x> <y> <radius>", file=sys.stderr); return
    x, y, r = int(args[0]), int(args[1]), int(args[2])
    result = api_get(f'/api/state/gas?x={x}&y={y}&radius={r}')
    if result.get("success"):
        d = result.get("data", {})
        gases = d.get("gases", [])
        total = sum(g.get('mass', 0) for g in gases)
        print(f"Gas analysis around ({x},{y}) radius {r}:")
        for g in sorted(gases, key=lambda x: x.get('mass', 0), reverse=True):
            pct = g['mass'] / total * 100 if total > 0 else 0
            print(f"  {g['gas']:20s} {g['mass']:8.1f} kg ({pct:4.1f}%) across {g.get('count', 0)} cells")
    else:
        print(f"Error: {result.get('errorMessage', result.get('error', 'unknown'))}", file=sys.stderr)

def cmd_registry(args):
    if not args:
        print("Usage: registry <buildings|elements|techs> [filter]", file=sys.stderr); return
    sub = args[0]
    filt = args[1] if len(args) > 1 else None
    if sub == "buildings":
        data = get_data('/api/registry/buildings')
    elif sub == "elements":
        data = get_data('/api/registry/elements')
    elif sub == "techs":
        data = get_data('/api/registry/techs')
    else:
        print(f"Unknown registry: {sub}", file=sys.stderr); return

    if not data: return
    if filt:
        filt_lower = filt.lower()
        data = [x for x in data if filt_lower in x.get('id', '').lower() or filt_lower in x.get('name', '').lower()]

    if isinstance(data, list):
        if sub == "buildings":
            for b in sorted(data, key=lambda x: x.get('id', '')):
                print(f"  {b.get('id', '?'):30s} {b.get('name', '?'):30s} {b.get('width',0)}x{b.get('height',0)}  rule={b.get('buildLocationRule','?')}  mat={b.get('materialCategory','?')}")
        elif sub == "elements":
            for e in sorted(data, key=lambda x: x.get('name', '')):
                print(f"  {e.get('name', '?'):20s} id={e.get('id', '?'):30s} state={e.get('state','?'):8s} SHC={e.get('specificHeatCapacity',0):.1f} TC={e.get('thermalConductivity',0):.2f}")
        elif sub == "techs":
            for t in sorted(data, key=lambda x: x.get('id', '')):
                reqs = ",".join(t.get('requiredTechs', []))
                print(f"  {t.get('id', '?'):30s} {t.get('name', '?'):30s} req={reqs}")
    else:
        print(json.dumps(data, indent=2, ensure_ascii=False))

def cmd_events(args):
    since = int(args[0]) if args else -1
    r = api_get(f'/api/state/events?since={since}')
    if r.get("success"):
        d = r.get("data", {})
        events = d.get("events", [])
        for e in events:
            print(f"  [{e.get('id',0):4d}] [{e.get('severity','?'):8s}] [{e.get('category','?')}] {e.get('title','')} {e.get('message','')}")
        if events:
            print(f"  --- next_seq={d.get('nextSeq', since)} (showing {len(events)} events) ---")
    else:
        print(f"Error: {r.get('errorMessage', r.get('error', 'unknown'))}", file=sys.stderr)

def cmd_dig(args):
    if len(args) < 4: print("Usage: dig <x> <y> <width> <height>", file=sys.stderr); return
    x, y, w, h = int(args[0]), int(args[1]), int(args[2]), int(args[3])
    post_action('/api/action/dig', {"x": x, "y": y, "width": w, "height": h})

def cmd_build(args):
    if len(args) < 3: print("Usage: build <buildingId> <x> <y>", file=sys.stderr); return
    bid, x, y = args[0], int(args[1]), int(args[2])
    r = post_action('/api/action/build', {"buildingId": bid, "x": x, "y": y})
    if r.get("success"):
        d = r.get("data", {})
        if d.get("materialInfo"):
            print("\nMaterial check:")
            for m in d["materialInfo"]:
                status = "ok" if m.get("available") else "missing"
                print(f"  [{status}] {m.get('category','?')}: {m.get('bestElement','none available')}")

def cmd_deconstruct(args):
    if len(args) < 2: print("Usage: deconstruct <x> <y>", file=sys.stderr); return
    x, y = int(args[0]), int(args[1])
    post_action('/api/action/deconstruct', {"x": x, "y": y})

def cmd_prioritize(args):
    if len(args) < 3: print("Usage: prioritize <x> <y> <priority 1-9>", file=sys.stderr); return
    x, y, p = int(args[0]), int(args[1]), int(args[2])
    post_action('/api/action/prioritize', {"x": x, "y": y, "priority": p})

def cmd_research_select(args):
    if not args: print("Usage: research <techId>", file=sys.stderr); return
    post_action('/api/action/research', {"techId": args[0]})

def cmd_mop(args):
    if len(args) < 2: print("Usage: mop <x> <y>", file=sys.stderr); return
    x, y = int(args[0]), int(args[1])
    post_action('/api/action/mop', {"x": x, "y": y})

def cmd_harvest(args):
    if len(args) < 2: print("Usage: harvest <x> <y>", file=sys.stderr); return
    x, y = int(args[0]), int(args[1])
    post_action('/api/action/harvest', {"x": x, "y": y})

def cmd_pause(args):
    reason = " ".join(args) if args else None
    r = api_post('/api/action/pause', {"reason": reason or "user request"})
    if r.get("success"): print("Paused")
    else: print(f"Error: {r.get('errorMessage', r.get('error'))}", file=sys.stderr)

def cmd_unpause(args):
    speed = int(args[0]) if args else 1
    r = api_post('/api/action/unpause', {"speed": speed})
    if r.get("success"): print(f"Unpaused at {speed}x")
    else: print(f"Error: {r.get('errorMessage', r.get('error'))}", file=sys.stderr)

def cmd_speed(args):
    if not args: print("Usage: speed <1|2|3>", file=sys.stderr); return
    s = int(args[0])
    r = api_post('/api/action/speed', {"speed": s})
    if r.get("success"): print(f"Speed set to {s}x")
    else: print(f"Error: {r.get('errorMessage', r.get('error'))}", file=sys.stderr)

def cmd_batch(args):
    if not args: print("Usage: batch <json_file>", file=sys.stderr); return
    with open(args[0]) as f:
        actions = json.load(f)
    if isinstance(actions, list):
        actions = {"actions": actions}
    r = api_post('/api/action/batch', actions)
    if r.get("success"):
        d = r.get("data", {})
        print(f"Batch: {d.get('successCount',0)} succeeded, {d.get('failCount',0)} failed of {d.get('total',0)}")
        for res in d.get("results", []):
            st = "✓" if res.get("success") else "✗"
            print(f"  {st} {res.get('type','?'):15s} {res.get('buildingId','')}")
    else:
        print(f"Error: {r.get('errorMessage', r.get('error'))}", file=sys.stderr)

def cmd_save(args):
    name = args[0] if args else None
    post_action('/api/action/save', {"name": name})

def cmd_load(args):
    if not args: print("Usage: load <save_name>", file=sys.stderr); return
    post_action('/api/action/load', {"name": args[0]})

def cmd_priority_global(args):
    if len(args) < 2: print("Usage: priority_global <target> <priority 1-9>", file=sys.stderr); return
    post_action('/api/action/priority_global', {"target": args[0], "priority": int(args[1])})

def cmd_priority_type(args):
    if len(args) < 2: print("Usage: priority_type <buildingType> <priority 1-9>", file=sys.stderr); return
    post_action('/api/action/priority_type', {"buildingType": args[0], "priority": int(args[1])})

def cmd_camera(args):
    if len(args) < 2: print("Usage: camera <x> <y> [zoom]", file=sys.stderr); return
    x, y = int(args[0]), int(args[1])
    body = {"x": x, "y": y}
    if len(args) > 2: body["zoom"] = float(args[2])
    post_action('/api/action/camera', body)

def cmd_explore(args):
    """AI-friendly area summary: combines cell data with building/dupe info."""
    if len(args) < 4: print("Usage: explore <x> <y> <width> <height>", file=sys.stderr); return
    x, y, w, h = int(args[0]), int(args[1]), int(args[2]), int(args[3])

    cells_r = api_get(f'/api/state/cells?x={x}&y={y}&width={w}&height={h}')
    buildings_r = api_get('/api/state/buildings')
    dups_r = api_get('/api/state/duplicants')

    print(f"\n=== Explore ({x},{y}) to ({x+w},{y+h}) ===")

    # Element distribution
    if cells_r.get("success"):
        cells = cells_r["data"].get("cells", [])
        dist = {}
        for c in cells:
            el = c.get("element", "Vacuum")
            if el not in dist: dist[el] = {"count": 0, "mass": 0, "state": "solid"}
            dist[el]["count"] += 1
            dist[el]["mass"] += c.get("massKg", 0)
            if c.get("isLiquid"): dist[el]["state"] = "liquid"
            elif c.get("isGas"): dist[el]["state"] = "gas"
            elif c.get("isVacuum"): dist[el]["state"] = "vacuum"

        print("\nElement distribution:")
        for el, info in sorted(dist.items(), key=lambda x: x[1]["count"], reverse=True):
            print(f"  {info['state']:8s} {el:20s} {info['count']:4d} cells {info['mass']:8.1f} kg")

        diggable = sum(1 for c in cells if c.get("isDiggable"))
        print(f"\nDiggable cells: {diggable}/{len(cells)}")

    # Buildings in area
    if buildings_r.get("success"):
        buildings = buildings_r.get("data", buildings_r)
        if isinstance(buildings, list):
            area_buildings = [b for b in buildings
                              if x <= b.get('x', -1) < x + w and y <= b.get('y', -1) < y + h]
            if area_buildings:
                print(f"\nBuildings in area:")
                for b in area_buildings:
                    print(f"  {b.get('id','?'):25s} at ({b.get('x')},{b.get('y')}) {b.get('width',1)}x{b.get('height',1)}")

    # Dupes in area
    if dups_r.get("success"):
        dups = dups_r.get("data", dups_r)
        if isinstance(dups, list):
            area_dupes = [d for d in dups
                          if x <= d.get('x', -1) < x + w and y <= d.get('y', -1) < y + h]
            if area_dupes:
                print(f"\nDuplicants in area:")
                for d in area_dupes:
                    print(f"  {d.get('name','?'):15s} at ({d.get('x')},{d.get('y')}) hp={d.get('health',0)}")

    print()

CMD_MAP = {
    "health": cmd_health, "status": cmd_status,
    "resources": cmd_resources, "buildings": cmd_buildings,
    "duplicants": cmd_duplicants, "research": cmd_research,
    "rooms": cmd_rooms,
    "cell": cmd_cell, "cells": cmd_cells, "slice": cmd_slice,
    "gas": cmd_gas, "explore": cmd_explore,
    "registry": cmd_registry,
    "events": cmd_events,
    "dig": cmd_dig, "build": cmd_build, "deconstruct": cmd_deconstruct,
    "prioritize": cmd_prioritize, "research_select": cmd_research_select,
    "mop": cmd_mop, "harvest": cmd_harvest,
    "pause": cmd_pause, "unpause": cmd_unpause, "speed": cmd_speed,
    "batch": cmd_batch, "camera": cmd_camera,
    "save": cmd_save, "load": cmd_load,
    "priority_global": cmd_priority_global, "priority_type": cmd_priority_type,
}

def main():
    if len(sys.argv) < 2 or sys.argv[1] in ("-h", "--help", "help"):
        print(__doc__)
        return

    cmd = sys.argv[1]
    args = sys.argv[2:]

    if cmd in CMD_MAP:
        CMD_MAP[cmd](args)
    else:
        print(f"Unknown command: {cmd}", file=sys.stderr)
        print("Available:", ", ".join(sorted(CMD_MAP.keys())), file=sys.stderr)
        sys.exit(1)

if __name__ == "__main__":
    main()