#!/usr/bin/env python
"""MCSJ Utility Balance Adjustment speed racer.

Runs one guarded balance-adjustment transaction locally without agent/chat pauses.
The timed interval starts at the first batch-id focus keystroke and ends when the
refreshed Utility Account Maintenance balance is semantically verified.
"""
from __future__ import annotations

import argparse, base64, json, socket, subprocess, sys, time, urllib.request
from decimal import Decimal
from pathlib import Path

HOST, PORT = "127.0.0.1", 9810
PANEL = "util.client.fx.EFXUtadjbatchPanel"
ACCOUNT_PANEL = "util.client.fx.EFXUtmastPanel"
PRINT_PANEL = "misc.components.fx.EFXPrtToScreenPanel"
ROOT = Path("C:/MCSJAgent")
FX = ROOT / "mcsj-fast-fx.py"
CUI = Path("C:/Users/Administrator/AppData/Local/Programs/Cua/cua-driver/bin/cua-driver.exe")
UIA_WORKER = ROOT / "mcsj-uia-worker.ps1"
JAB_CACHE = ROOT / "cache/mcsj-jab-cache.json"
PROOF = ROOT / "proof"


def token(value: str) -> str:
    if value == "": return "."
    return base64.urlsafe_b64encode(value.encode()).decode().rstrip("=")


def agent(command: str, *, modal: bool = False, timeout: float = 3.0) -> dict:
    with socket.create_connection((HOST, PORT), timeout=timeout) as s:
        s.settimeout(timeout)
        s.sendall((command + "\n").encode())
        s.shutdown(socket.SHUT_WR)
        if modal:
            return {"status": "DISPATCHED", "command": command.split()[0]}
        chunks = []
        while True:
            data = s.recv(65536)
            if not data: break
            chunks.append(data)
    raw = b"".join(chunks).decode(errors="replace")
    obj = json.loads(raw)
    if obj.get("status") != "OK": raise RuntimeError(f"Agent rejected {command.split()[0]}: {obj}")
    return obj


def fxstate() -> dict:
    return agent("FXSTATE", timeout=5)


def fire(panel: str, text: str, *, modal: bool = False) -> dict:
    return agent(f"FXFIREBUTTON {token(panel)} {token(text)} false", modal=modal)


def select_radio(text: str) -> dict:
    return agent(f"FXSELECTRADIO {token(PANEL)} {token(text)} false true")


def edit_currency(expected: str, new: str) -> dict:
    cmd = "FXEDITTEXT {} 0 {} {} * 0 {} {}".format(
        token(PANEL), token("misc.components.fx.EFXCurrencyField"), token("EFXTextField"),
        token(expected), token(new))
    return agent(cmd)


def jab(method: str, path: str, payload=None) -> dict:
    data = json.dumps(payload).encode() if payload is not None else None
    req = urllib.request.Request("http://127.0.0.1:8765" + path, data=data,
        headers={"Content-Type":"application/json"} if data else {}, method=method)
    with urllib.request.urlopen(req, timeout=10) as r:
        return json.loads(r.read().decode())


def cua(tool: str, payload: dict, timeout: float = 6.0) -> dict:
    p = subprocess.run([str(CUI), "call", tool, "--json"], input=json.dumps(payload),
        capture_output=True, text=True, timeout=timeout)
    if p.returncode != 0: raise RuntimeError(f"CUA {tool} failed: {p.stderr} {p.stdout}")
    return json.loads(p.stdout) if p.stdout.strip().startswith("{") else {"status":"OK","raw":p.stdout.strip()}


def press(pid: int, hwnd: int, key: str) -> None:
    cua("press_key", {"pid":pid,"window_id":hwnd,"key":key,"dispatch":"background"})


def type_text(pid: int, hwnd: int, text: str) -> None:
    cua("type_text", {"pid":pid,"window_id":hwnd,"text":text,"delay_ms":0,"dispatch":"background"}, 8)


class UiaWorker:
    def __init__(self, pid: int):
        self.pid = pid
        self.p = subprocess.Popen(["powershell.exe","-NoProfile","-ExecutionPolicy","Bypass","-File",str(UIA_WORKER)],
            stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, bufsize=1)
    def request(self, action: str, window: str, button: str | None = None) -> dict:
        req={"action":action,"pid":self.pid,"window":window}
        if button is not None: req["button"]=button
        self.p.stdin.write(json.dumps(req,separators=(",",":"))+"\n"); self.p.stdin.flush()
        line=self.p.stdout.readline()
        if not line: raise RuntimeError("UIA worker exited: "+self.p.stderr.read())
        return json.loads(line)
    def snapshot(self, title: str) -> dict: return self.request("snapshot",title)
    def invoke(self, title: str, button: str) -> dict:
        out=self.request("invoke",title,button)
        if out.get("status")!="OK": raise RuntimeError(f"UIA invoke failed: {out}")
        return out
    def wait(self, title: str, timeout: float, contains: list[str] | None = None) -> dict:
        deadline=time.perf_counter()+timeout
        while time.perf_counter()<deadline:
            out=self.snapshot(title)
            if out.get("status")=="OK":
                hay="\n".join(str(i.get("name") or "")+"\n"+str(i.get("value") or "") for i in out.get("items",[]))
                if not contains or all(x in hay for x in contains): return out
            time.sleep(.03)
        raise TimeoutError(f"window {title!r} missing expected text {contains}")
    def optional(self,title: str, timeout: float=.6) -> dict | None:
        deadline=time.perf_counter()+timeout
        while time.perf_counter()<deadline:
            out=self.snapshot(title)
            if out.get("status")=="OK": return out
            time.sleep(.03)
        return None
    def close(self):
        try: self.p.stdin.close(); self.p.terminate()
        except Exception: pass


def panel(state: dict, class_name: str) -> dict:
    matches=[p for p in state.get("panels",[]) if p.get("panelClass")==class_name]
    if len(matches)!=1: raise RuntimeError(f"expected one {class_name}, got {len(matches)}")
    return matches[0]


def find_account_menu_id() -> int:
    obj=json.loads(JAB_CACHE.read_text(encoding="utf-8")); rows=obj["rows"]; by={r["id"]:r for r in rows}
    hits=[]
    for r in rows:
        if r.get("name")!="Account Maintenance" or r.get("role")!="menu item": continue
        names=[]; n=r
        while n:
            names.append(n.get("name","")); n=by.get(n.get("parent"))
        if "Utility Accounts" in names and "Utility/Property Tax Billing" in names: hits.append(r["id"])
    if len(hits)!=1: raise RuntimeError(f"exact Utility Account Maintenance menu match count={len(hits)}")
    return hits[0]


def wait_fx(predicate, timeout: float, label: str) -> dict:
    deadline=time.perf_counter()+timeout; last=None
    while time.perf_counter()<deadline:
        try:
            last=fxstate()
            if predicate(last): return last
        except (OSError, json.JSONDecodeError, RuntimeError): pass
        time.sleep(.04)
    raise TimeoutError(f"FX wait timed out: {label}; last={last}")


def ensure_preflight(expected_db: str) -> tuple[int,int,int]:
    p=subprocess.run([sys.executable,str(FX),"ensure"],capture_output=True,text=True,timeout=40)
    if p.returncode: raise RuntimeError(p.stderr+p.stdout)
    ensured=json.loads(p.stdout.strip().splitlines()[-1])
    if expected_db.lower() not in ensured.get("databaseTitle","").lower(): raise RuntimeError(f"wrong database: {ensured}")
    js=jab("GET","/state"); pid=int(js["pid"]); hwnd=int(js["hwnd"])
    if pid!=int(ensured["pid"]): raise RuntimeError("JAB/Agent PID mismatch")
    st=fxstate(); psetup=panel(st,PANEL)
    if psetup.get("tables"): raise RuntimeError("adjustment entry grid is open; require blank setup screen")
    batch=[c for c in psetup["controls"] if c.get("class","").endswith("EFXUpperMaskField")]
    if len(batch)!=1 or batch[0].get("text")!="": raise RuntimeError(f"batch setup is not blank: {batch}")
    focused=[c for c in psetup["controls"] if c.get("focused")]
    if focused: raise RuntimeError(f"setup focus must be empty for deterministic first Tab: {focused}")
    return pid,hwnd,find_account_menu_id()


def run(args) -> dict:
    amount=Decimal(args.principal); before=Decimal(args.expected_before); expected_after=before+amount
    report_amount=(f"{abs(amount):.2f}-" if amount < 0 else f"{amount:.2f}")
    pid,hwnd,account_menu_id=ensure_preflight(args.database)
    if not args.execute:
        return {"status":"PREFLIGHT_OK","pid":pid,"hwnd":hwnd,"database":args.database,
            "batch":args.batch,"account":f"{args.account}-{args.sub}","expectedAfter":f"{expected_after:.2f}","menuId":account_menu_id}
    worker=UiaWorker(pid); marks={}; start=time.perf_counter()
    def mark(name): marks[name]=round(time.perf_counter()-start,3)
    try:
        # Setup and row entry: one deterministic keyboard route, exact semantic buttons.
        press(pid,hwnd,"tab"); type_text(pid,hwnd,args.batch); press(pid,hwnd,"tab")
        fire(PANEL,"_Next"); fire(PANEL,"_Add")
        type_text(pid,hwnd,args.account); press(pid,hwnd,"tab"); type_text(pid,hwnd,args.sub); press(pid,hwnd,"tab")
        press(pid,hwnd,"tab"); press(pid,hwnd,"tab"); type_text(pid,hwnd,args.code); press(pid,hwnd,"tab"); press(pid,hwnd,"tab")
        type_text(pid,hwnd,str(args.year)); press(pid,hwnd,"tab"); type_text(pid,hwnd,str(args.period)); press(pid,hwnd,"tab")
        edit_currency(".00",args.principal); press(pid,hwnd,"tab"); press(pid,hwnd,"tab"); type_text(pid,hwnd,args.description); press(pid,hwnd,"tab")
        mark("entry_complete")

        fire(PANEL,"_Save",modal=True)
        totals=worker.wait("Batch Totals",3,[args.batch,"1",args.principal])
        text="\n".join(str(i.get("name") or "")+str(i.get("value") or "") for i in totals["items"])
        if text.count(args.principal)<2: raise RuntimeError("Batch Totals principal/total did not both match")
        worker.invoke("Batch Totals","OK")
        warning=worker.optional("Message",.8)
        if warning:
            wtext="\n".join(str(i.get("name") or "") for i in warning["items"])
            if "Adjustment Date is not in the Current Year" not in wtext: raise RuntimeError("unexpected Message dialog: "+wtext)
            worker.invoke("Message","OK")
        saved=wait_fx(lambda s: panel(s,PANEL).get("tables") and panel(s,PANEL)["tables"][0]["rows"][0][0]!="INS",3,"saved row")
        row=panel(saved,PANEL)["tables"][0]["rows"][0]
        required=[args.account,args.sub,args.service,args.code,str(args.year),str(args.period),args.principal,"0.00",args.description]
        if not all(x in row for x in required): raise RuntimeError(f"saved row mismatch: {row}")
        mark("saved")

        fire(PANEL,"_Close"); select_radio("Batch Verification Listing"); fire(PANEL,"_Next",modal=True)
        worker.wait("Select Printing Method",3,["Select Printing Method:"])
        worker.invoke("Select Printing Method","Screen")
        def verification_ready(state):
            matches=[p for p in state.get("panels",[]) if p.get("panelClass")==PRINT_PANEL]
            if len(matches)!=1: return False
            text="\n".join(str(cell) for t in matches[0].get("tables",[]) for rowx in t.get("rows",[]) for cell in rowx)
            return args.batch in text and "Grand Totals:" in text
        report=wait_fx(verification_ready,5,"populated verification report")
        rp=panel(report,PRINT_PANEL); report_text="\n".join(str(cell) for t in rp.get("tables",[]) for rowx in t.get("rows",[]) for cell in rowx)
        for needed in [args.batch,f"{args.account}-{args.sub}",args.description,"Grand Totals:",report_amount]:
            if needed not in report_text: raise RuntimeError(f"verification report missing {needed!r}")
        fire(PRINT_PANEL,"_Close"); mark("verified")

        select_radio("Update Batch"); fire(PANEL,"_Next",modal=True)
        worker.wait("Select an Option",3,["updates journal entries","verified"]); worker.invoke("Select an Option","OK")
        pr=worker.wait("Print",6)
        upd=worker.wait("Batch Update",3,[args.batch,"Updating Masters"])
        upd_text="\n".join(str(i.get("name") or "")+" "+str(i.get("value") or "") for i in upd["items"])
        if upd_text.count("1")<2: raise RuntimeError("Batch Update counters did not show 1/1")
        press(pid,int(pr["hwnd"]),"escape")
        worker.wait("Batch Update",10,["Update Completed"]); worker.invoke("Batch Update","OK"); mark("updated")

        opened=jab("POST","/click",{"id":account_menu_id})
        if not opened.get("ok"): raise RuntimeError(f"account menu click failed: {opened}")
        wait_fx(lambda s:any(p.get("panelClass")==ACCOUNT_PANEL for p in s.get("panels",[])),3,"account panel")
        type_text(pid,hwnd,args.account); press(pid,hwnd,"tab"); type_text(pid,hwnd,args.sub); press(pid,hwnd,"tab")
        wait_fx(lambda s:any(c.get("text","").strip()==args.expected_owner for c in panel(s,ACCOUNT_PANEL)["controls"]),3,"account load")
        for _ in range(7): press(pid,hwnd,"f4")
        final=wait_fx(lambda s:any(t.get("rows") and any(r[0]==args.service for r in t["rows"]) for t in panel(s,ACCOUNT_PANEL).get("tables",[])),4,f"{args.service} balance")
        ap=panel(final,ACCOUNT_PANEL); bt=next(t for t in ap["tables"] if t.get("rows") and any(r[0]==args.service for r in t["rows"]))
        water=next((r for r in bt["rows"] if r[0]==args.service),None)
        if water is None: raise RuntimeError(f"service row not found: {args.service}")
        actual=Decimal(water[2])
        if actual!=expected_after: raise RuntimeError(f"final balance mismatch actual={actual} expected={expected_after}")
        elapsed=round(time.perf_counter()-start,3); mark("after_verified")
        PROOF.mkdir(parents=True,exist_ok=True)
        shot=PROOF/f"speed-racer-{args.batch}-after.png"
        subprocess.run(["powershell.exe","-NoProfile","-ExecutionPolicy","Bypass","-File",str(ROOT/"printwindow-shot.ps1"),"-Path",str(shot)],capture_output=True,text=True,timeout=20,check=True)
        fire(ACCOUNT_PANEL,"_Close")
        result={"status":"UPDATED_AND_AFTER_VERIFIED","database":args.database,"batch":args.batch,
            "account":f"{args.account}-{args.sub}","service":args.service,"principal":args.principal,"before":f"{before:.2f}",
            "after":f"{actual:.2f}","elapsedSeconds":elapsed,"sub10":elapsed<10,"milestones":marks,"proof":str(shot)}
        (PROOF/f"speed-racer-{args.batch}.json").write_text(json.dumps(result,indent=2),encoding="utf-8")
        return result
    finally:
        worker.close()


def main() -> int:
    ap=argparse.ArgumentParser()
    ap.add_argument("--execute",action="store_true",help="perform the mutation; otherwise preflight only")
    ap.add_argument("--database",default="chestervillagenydb")
    ap.add_argument("--batch",required=True); ap.add_argument("--account",default="601"); ap.add_argument("--sub",default="2")
    ap.add_argument("--service",default="Water"); ap.add_argument("--expected-owner",default="CHESTER ACQUISITIONS LLC")
    ap.add_argument("--code",default="BAL"); ap.add_argument("--year",type=int,default=2026); ap.add_argument("--period",type=int,default=1)
    ap.add_argument("--principal",default="-1.00"); ap.add_argument("--description",required=True); ap.add_argument("--expected-before",required=True)
    args=ap.parse_args()
    try:
        out=run(args); print(json.dumps(out,indent=2)); return 0
    except Exception as e:
        print(json.dumps({"status":"ERROR","stageError":str(e)},indent=2),file=sys.stderr); return 1

if __name__=="__main__": raise SystemExit(main())