aboutsummaryrefslogtreecommitdiffstats
path: root/server/patch_server.py
blob: c9d27051867b2854a31f1d320dd32719d3bcf543 (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
1466
1467
1468
1469
1470
1471
1472
1473
1474
1475
1476
1477
1478
1479
1480
1481
1482
1483
1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
1497
1498
1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543
1544
1545
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
1556
1557
1558
1559
1560
1561
1562
1563
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
1574
1575
1576
1577
1578
1579
1580
1581
1582
1583
1584
1585
1586
1587
1588
1589
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
1600
1601
1602
1603
1604
1605
1606
1607
1608
1609
1610
1611
1612
1613
1614
1615
1616
1617
1618
1619
1620
1621
1622
1623
1624
1625
1626
1627
1628
1629
1630
1631
1632
1633
1634
1635
1636
1637
1638
1639
1640
1641
1642
1643
1644
1645
1646
1647
1648
1649
1650
1651
1652
1653
1654
1655
1656
1657
1658
1659
1660
1661
1662
1663
1664
1665
1666
1667
1668
1669
1670
1671
1672
1673
1674
1675
1676
1677
1678
1679
1680
1681
1682
1683
1684
1685
1686
1687
1688
1689
1690
1691
1692
1693
1694
1695
1696
1697
1698
1699
1700
1701
1702
1703
1704
1705
1706
1707
1708
1709
1710
1711
1712
1713
1714
1715
1716
1717
1718
1719
1720
1721
1722
1723
1724
1725
1726
1727
1728
1729
1730
1731
1732
1733
1734
1735
1736
1737
1738
1739
1740
1741
1742
1743
1744
1745
1746
1747
1748
1749
1750
1751
1752
1753
1754
1755
1756
1757
1758
1759
1760
1761
1762
1763
1764
1765
1766
1767
1768
1769
1770
1771
1772
1773
1774
1775
1776
1777
1778
1779
1780
1781
1782
1783
1784
1785
1786
1787
1788
1789
1790
1791
1792
1793
1794
1795
1796
1797
1798
1799
1800
1801
1802
1803
1804
1805
1806
1807
1808
1809
1810
1811
1812
1813
1814
1815
1816
1817
1818
1819
1820
1821
1822
1823
1824
1825
1826
1827
1828
1829
1830
1831
1832
1833
1834
1835
1836
1837
1838
1839
1840
1841
1842
1843
1844
1845
1846
1847
1848
1849
1850
1851
1852
1853
1854
1855
1856
1857
1858
1859
1860
1861
1862
1863
1864
1865
1866
1867
1868
1869
1870
1871
1872
1873
1874
1875
1876
1877
1878
1879
1880
1881
1882
1883
1884
1885
1886
1887
1888
1889
1890
1891
1892
1893
1894
1895
1896
1897
1898
1899
1900
1901
1902
1903
1904
1905
1906
1907
1908
1909
1910
1911
1912
1913
1914
1915
1916
1917
1918
1919
1920
1921
1922
1923
1924
1925
1926
1927
1928
1929
1930
1931
1932
1933
1934
1935
1936
1937
1938
1939
1940
1941
1942
1943
1944
1945
1946
1947
1948
1949
1950
1951
1952
1953
1954
1955
1956
1957
1958
1959
1960
1961
1962
1963
1964
1965
1966
1967
1968
1969
1970
1971
1972
1973
1974
1975
1976
1977
1978
1979
1980
1981
1982
1983
1984
1985
1986
1987
1988
1989
1990
1991
1992
1993
1994
1995
1996
1997
1998
1999
2000
2001
2002
2003
2004
2005
2006
2007
2008
2009
#!/usr/bin/env python3
"""Inject idatui's extra ida-pro-mcp tools into the installed server package.

DEPRECATED along with the ida-pro-mcp transport: the default backend is now the
idalib worker (idatui/worker.py), which registers these same tools in-process and
needs no patching. Kept only for `--backend mcp`; slated for removal.

ida-pro-mcp lacks a few tools idatui needs. Rather than vendor/fork the server,
we keep the tool source here and inject it (idempotently) into the installed
``api_types.py``. That module is imported by every worker
(``python -m ida_pro_mcp.idalib_server``), so the tools register themselves via
``@tool`` on the shared ``MCP_SERVER`` — no server code is forked, and re-running
this (spawn.sh does, on every start) re-applies it after a reinstall/upgrade.

Injected tools:
  * ``del_type``      — delete a named local type (struct editor CRUD).
  * ``func_types``    — structured decompiler types for a function (prototype +
                        local variables), so clients don't parse pseudocode text.
  * ``set_lvar_type`` — set a decompiler local variable's type; works on auto/
                        register vars too (the stock set_type only updates lvars
                        that already have user-saved info).

The block between the BEGIN/END markers is *replaced* on each run, so editing
BODY here and restarting the supervisor updates the tools.

Run with the *same* interpreter the server uses (the idalib-mcp entry point's
``/usr/bin/python``), so it patches the file the workers actually import.
Changing a tool needs a supervisor restart so workers respawn.
"""
from __future__ import annotations

import importlib.util
import pathlib
import sys

BEGIN = "# >>> idatui-ext: begin (auto-injected by server/patch_server.py) >>>"
END = "# <<< idatui-ext: end <<<"

# Appended to ida_pro_mcp/ida_mcp/api_types.py, which already imports
# ``Annotated``, ``tool``, ``idasync``, ``ida_typeinf``, ``parse_address`` and
# ``_parse_type_tinfo``.
BODY = '''
def _idatui_lv_get(x):
    return x() if callable(x) else x


@tool
@idasync
def resolve_names(
    queries: Annotated[list, "Symbol name(s) to resolve to their OWN address"],
) -> list:
    """Resolve named locations (functions, labels like loc_/locret_, data) to the
    exact address the NAME denotes, via get_name_ea. Unlike lookup_funcs, a
    mid-function label resolves to the label's address, not the containing
    function's entry."""
    import idaapi
    qs = queries if isinstance(queries, list) else [queries]
    out = []
    for q in qs:
        q = str(q).strip()
        ea = idaapi.get_name_ea(idaapi.BADADDR, q)
        out.append({"query": q, "ea": (hex(ea) if ea != idaapi.BADADDR else None)})
    return out


@tool
@idasync
def del_type(
    name: Annotated[str, "Local type name to delete (struct/union/enum/typedef)"],
) -> dict:
    """Delete a named local type from the local type library."""
    til = ida_typeinf.get_idati()
    ok = ida_typeinf.del_named_type(til, name, ida_typeinf.NTF_TYPE)
    if not ok:
        return {"name": name, "error": f"Type '{name}' not found or could not be deleted"}
    return {"name": name, "deleted": True}


@tool
@idasync
def func_types(
    addr: Annotated[str, "Function address or name"],
) -> dict:
    """Structured decompiler types for a function: its prototype plus each local
    variable (name/type/is_arg). Lets clients read/edit types without parsing
    pseudocode text."""
    import ida_hexrays
    import idaapi

    def _tstr(tif):
        try:
            s = tif.dstr()
            if s:
                return s
        except Exception:
            pass
        return str(tif)

    ea = parse_address(addr)
    f = idaapi.get_func(ea)
    if not f:
        return {"addr": str(addr), "error": "no function at address"}
    try:
        cf = ida_hexrays.decompile(f.start_ea)
    except Exception as e:
        return {"addr": hex(f.start_ea), "error": f"decompile failed: {e}"}
    if cf is None:
        return {"addr": hex(f.start_ea), "error": "decompilation failed"}
    name = idaapi.get_func_name(f.start_ea) or ""
    try:
        proto = ida_typeinf.print_tinfo(
            "", 0, 0, ida_typeinf.PRTYPE_1LINE, cf.type, name, "")
    except Exception:
        proto = ""
    lvars = []
    for lv in cf.get_lvars():
        try:
            ty = _tstr(_idatui_lv_get(lv.type))
        except Exception:
            ty = ""
        lvars.append({
            "name": _idatui_lv_get(lv.name),
            "type": ty,
            "is_arg": bool(_idatui_lv_get(lv.is_arg_var)),
        })
    return {
        "addr": hex(f.start_ea),
        "name": name,
        "prototype": (proto or "").strip(),
        "lvars": lvars,
    }


@tool
@idasync
def set_lvar_type(
    addr: Annotated[str, "Function address or name"],
    variable: Annotated[str, "Local variable name"],
    type: Annotated[str, "New C type for the variable"],
) -> dict:
    """Set a decompiler local variable's type. Handles auto/register vars (unlike
    set_type, which only updates lvars that already have user-saved info)."""
    import ida_hexrays
    import idaapi

    ea = parse_address(addr)
    f = idaapi.get_func(ea)
    if not f:
        return {"error": "no function at address"}
    try:
        cf = ida_hexrays.decompile(f.start_ea)
    except Exception as e:
        return {"error": f"decompile failed: {e}"}
    if cf is None:
        return {"error": "decompilation failed"}
    target = None
    for lv in cf.get_lvars():
        if _idatui_lv_get(lv.name) == variable:
            target = lv
            break
    if target is None:
        return {"error": f"local variable {variable!r} not found"}
    try:
        tif = _parse_type_tinfo(type)
    except Exception as e:
        return {"error": f"bad type {type!r}: {e}"}
    lsi = ida_hexrays.lvar_saved_info_t()
    try:
        lsi.ll = target
    except Exception:
        try:
            lsi.ll.location = _idatui_lv_get(target.location)
            lsi.ll.defea = target.defea
        except Exception as e:
            return {"error": f"could not locate variable: {e}"}
    lsi.type = tif
    ok = bool(ida_hexrays.modify_user_lvar_info(
        f.start_ea, ida_hexrays.MLI_TYPE, lsi))
    return {"addr": hex(f.start_ea), "variable": variable, "type": type, "ok": ok}


@tool
@idasync
def file_regions() -> dict:
    """Loaded segments mapped to their raw file offsets (get_fileregion_offset),
    so clients can convert a virtual address to an on-disk file offset without a
    format-specific header parser. file_off is -1 for non-file-backed segments
    (e.g. .bss)."""
    import ida_segment
    import idaapi

    out = []
    seg = ida_segment.get_first_seg()
    while seg is not None:
        try:
            fo = int(idaapi.get_fileregion_offset(seg.start_ea))
        except Exception:
            fo = -1
        if fo < 0 or fo >= (1 << 48):
            fo = -1
        try:
            nm = ida_segment.get_segm_name(seg) or ""
        except Exception:
            nm = ""
        out.append({"start": hex(seg.start_ea), "end": hex(seg.end_ea),
                    "file_off": fo, "name": nm})
        seg = ida_segment.get_next_seg(seg.start_ea)
    return {"regions": out}


@tool
@idasync
def make_string(
    addr: Annotated[str, "Address of the string start"],
    length: Annotated[int, "Length in bytes (0 = auto-detect to the terminator)"] = 0,
    kind: Annotated[str, "String kind: c | c16 | c32 | pascal"] = "c",
) -> dict:
    """Create a string literal at ``addr`` (IDA's 'A'). ``length`` 0 auto-detects
    to the terminator. Undefines any items in the way first, like the UI does.
    Returns the created byte size and the decoded contents."""
    import ida_bytes
    import ida_nalt

    ea = parse_address(addr)
    strtype = {
        "c": ida_nalt.STRTYPE_C,
        "c16": ida_nalt.STRTYPE_C_16,
        "c32": ida_nalt.STRTYPE_C_32,
        "pascal": ida_nalt.STRTYPE_PASCAL,
    }.get(str(kind).lower(), ida_nalt.STRTYPE_C)
    n = max(int(length), 0)
    # Free any existing item(s) so create_strlit can carve the literal.
    ida_bytes.del_items(ea, ida_bytes.DELIT_SIMPLE, n if n > 0 else 1)
    ok = bool(ida_bytes.create_strlit(ea, n, strtype))
    if not ok:
        return {"addr": addr, "ok": False, "error": "create_strlit failed"}
    size = int(ida_bytes.get_item_size(ea))
    try:
        raw = ida_bytes.get_strlit_contents(ea, -1, strtype)
        text = raw.decode("utf-8", "replace") if raw else ""
    except Exception:
        text = ""
    return {"addr": addr, "ok": True, "size": size, "text": text}


@tool
@idasync
def read_raw(
    addr: Annotated[str, "Start address (hex or name)"],
    size: Annotated[int, "Number of bytes to read"],
) -> dict:
    """Read ``size`` bytes at ``addr`` as ONE contiguous lowercase hex string
    (no per-byte '0x'/spaces). The hot path for the hex view and disasm opcode
    bytes.

    Fast: does a single bulk ``ida_bytes.get_bytes`` (C-speed) instead of the
    per-byte read_bytes_bss_safe loop (2 IDA calls/byte). Unloaded bytes come
    back from IDA as the 0xFF sentinel, so we only re-check is_loaded for the
    (usually sparse) 0xFF bytes and zero the genuinely-unloaded ones — matching
    get_bytes' bss semantics without paying per-byte for the whole range.

    Encoding is compact hex (~2.5x smaller than get_bytes' '0x..'-with-spaces)
    and, unlike get_bytes, does not truncate on large reads."""
    import ida_bytes

    ea = parse_address(addr)
    n = max(int(size), 0)
    if n == 0:
        return {"addr": addr, "hex": "", "n": 0}
    raw = ida_bytes.get_bytes(ea, n)
    if raw is None or len(raw) < n:  # nothing (or not all) mapped
        base = bytearray(raw or b"")
        base.extend(b"\\xff" * (n - len(base)))
        raw = bytes(base)
    ba = bytearray(raw)
    # Only unloaded bytes read as 0xFF; correct just those to 0 (bss => zero).
    i = ba.find(0xFF)
    while i != -1:
        if not ida_bytes.is_loaded(ea + i):
            ba[i] = 0
        i = ba.find(0xFF, i + 1)
    return {"addr": addr, "hex": bytes(ba).hex(), "n": len(ba)}


def _idatui_head_row(ea):
    """One flat-listing row for the head at ``ea``: kind (code/data/unknown),
    byte size, rendered text, and any symbol name."""
    import ida_bytes
    import ida_lines
    import ida_name

    f = ida_bytes.get_flags(ea)
    if ida_bytes.is_code(f):
        kind = "code"
    elif ida_bytes.is_data(f):
        kind = "data"
    else:
        kind = "unknown"
    line = ida_lines.generate_disasm_line(ea, 0)
    text = ida_lines.tag_remove(line) if line else ""
    text = " ".join(text.split())  # collapse IDA's column padding
    row = {
        "ea": hex(ea),
        "kind": kind,
        "size": int(ida_bytes.get_item_size(ea)),
        "text": text,
    }
    if line:
        # Keep IDA's own token classification for syntax highlighting. Built from
        # the SAME line as `text`, then whitespace-collapsed identically so the
        # two never disagree about what the row says.
        spans, ops = _idatui_spans(line)
        joined = "".join(t for _k, t in spans)
        if " ".join(joined.split()) == text:
            row["spans"] = spans
            # Where each operand sits in `text`. Comes out of the same tag walk
            # (free), and is what lets the client show WHICH literal a keypress
            # would reformat before you press it.
            if ops:
                row["ops"] = ops
    nm = ida_name.get_ea_name(ea)
    if nm:
        row["name"] = nm
    return row


#: IDA colour tag -> the semantic kind the TUI styles. IDA already classifies
#: every token in a disassembly line, for every processor it supports, so there
#: is nothing to lex: generate_disasm_line emits \x01<tag>text\x02<tag> and the
#: tag says what the text IS. A pygments assembly lexer would be a worse guess at
#: this and would need one dialect per architecture.
_IDATUI_SPAN_KINDS = {
    "insn": ("SCOLOR_INSN", "SCOLOR_KEYWORD", "SCOLOR_ASMDIR", "SCOLOR_MACRO"),
    "reg": ("SCOLOR_REG",),
    "num": ("SCOLOR_NUMBER", "SCOLOR_CHAR", "SCOLOR_BINPREF"),
    "str": ("SCOLOR_STRING",),
    # NB the real constant names: DATNAME/CODNAME, not "DNAME". Guessing here
    # fails silently — an unmapped tag renders as plain body text, so symbols
    # just quietly aren't blue and nothing tells you why.
    "name": ("SCOLOR_DATNAME", "SCOLOR_CODNAME", "SCOLOR_LOCNAME",
             "SCOLOR_IMPNAME", "SCOLOR_DEMNAME", "SCOLOR_LIBNAME",
             "SCOLOR_CNAME", "SCOLOR_DNAME",
             "SCOLOR_CREF", "SCOLOR_DREF", "SCOLOR_CREFTAIL", "SCOLOR_DREFTAIL"),
    "seg": ("SCOLOR_SEGNAME",),
    "cmt": ("SCOLOR_AUTOCMT", "SCOLOR_REGCMT", "SCOLOR_RPTCMT", "SCOLOR_VOIDOP"),
    "punct": ("SCOLOR_SYMBOL", "SCOLOR_ALTOP", "SCOLOR_HIDNAME"),
    "err": ("SCOLOR_ERROR",),
}


def _idatui_tag_map():
    """{tag character: kind}, built once from whatever this IDA actually has."""
    import ida_lines
    out = {}
    for kind, names in _IDATUI_SPAN_KINDS.items():
        for n in names:
            v = getattr(ida_lines, n, None)
            if isinstance(v, str) and v:
                out[v[0]] = kind
            elif isinstance(v, int):
                out[chr(v)] = kind
    return out


_IDATUI_TAGS = None
_IDATUI_OPND_TAGS = None


def _idatui_opnd_tag_map():
    """{tag character: operand index}. IDA wraps each operand of a disassembly
    line in COLOR_OPND1..8, so the line already says where operand N starts and
    ends -- no need to re-render operands with print_operand to find out (and
    the two agree exactly; checked over thousands of instructions)."""
    import ida_lines
    out = {}
    for i in range(1, 9):
        v = getattr(ida_lines, "COLOR_OPND%d" % i, None)
        if isinstance(v, int):
            out[chr(v)] = i - 1
        elif isinstance(v, str) and v:
            out[v[0]] = i - 1
    return out


def _idatui_spans(line):
    """(spans, ops) for a tagged disasm line.

    ``spans`` is [[kind, text], ...] with colour tags resolved; ``ops`` is
    [[start, end, n], ...], the extent of each operand in the SAME (collapsed)
    coordinates the row's ``text`` uses -- which is what lets a cursor column
    name the operand it is standing on.

    Unknown tags become 'text' rather than being dropped: a processor module can
    emit a colour we don't classify, and losing the characters would corrupt the
    line."""
    global _IDATUI_TAGS, _IDATUI_OPND_TAGS
    import ida_lines
    if _IDATUI_TAGS is None:
        _IDATUI_TAGS = _idatui_tag_map()
    if _IDATUI_OPND_TAGS is None:
        _IDATUI_OPND_TAGS = _idatui_opnd_tag_map()
    on, off, esc = "\x01", "\x02", "\x03"
    addr_tag = chr(getattr(ida_lines, "COLOR_ADDR", 0x28))
    addr_len = int(getattr(ida_lines, "COLOR_ADDR_SIZE", 16))
    spans, stack, buf = [], [], []   # stack entries: (kind, operand index|None)
    i, n = 0, len(line)

    def _opnd():
        for _k, o in reversed(stack):
            if o is not None:
                return o
        return None

    def flush():
        if buf:
            spans.append([stack[-1][0] if stack else "text", "".join(buf),
                          _opnd()])
            del buf[:]

    while i < n:
        ch = line[i]
        if ch == on and i + 1 < n:
            tag = line[i + 1]
            if tag == addr_tag:
                # An embedded target address, not display text: 16 hex digits
                # that must not reach the screen.
                i += 2 + addr_len
                continue
            flush()
            stack.append((_IDATUI_TAGS.get(tag, "text"),
                          _IDATUI_OPND_TAGS.get(tag)))
            i += 2
            continue
        if ch == off and i + 1 < n:
            flush()
            if stack:
                stack.pop()
            i += 2
            continue
        if ch == esc and i + 1 < n:      # escaped literal
            buf.append(line[i + 1])
            i += 2
            continue
        buf.append(ch)
        i += 1
    flush()
    # Collapse IDA's column padding EXACTLY as the plain text does. A run of
    # spaces can straddle two spans, so this walks characters rather than
    # collapsing each span on its own — otherwise the spans and `text` disagree
    # about the line and the row silently loses its highlighting.
    out, prev_space = [], False
    for kind, txt, opnd in spans:
        acc = []
        for ch in txt:
            if ch.isspace():
                if prev_space:
                    continue
                acc.append(" ")
                prev_space = True
            else:
                acc.append(ch)
                prev_space = False
        if acc:
            out.append([kind, "".join(acc), opnd])
    while out and out[0][1] == " ":
        out.pop(0)
    while out and out[-1][1] == " ":
        out.pop()
    if out and out[0][1].startswith(" "):
        out[0][1] = out[0][1].lstrip()
    if out and out[-1][1].endswith(" "):
        out[-1][1] = out[-1][1].rstrip()
    out = [s for s in out if s[1]]
    # Operand extents, in the coordinates of the collapsed text these spans
    # spell out. Adjacent spans of the same operand merge, so an operand like
    # ``[rbp+var_40]`` (five differently-coloured tokens) comes back as ONE
    # range -- which is the thing a cursor is inside of, and the thing a format
    # change applies to.
    ops, pos, cur, start = [], 0, None, 0
    for _kind, txt, opnd in out:
        if opnd != cur:
            if cur is not None and pos > start:
                ops.append([start, pos, cur])
            cur, start = opnd, pos
        pos += len(txt)
    if cur is not None and pos > start:
        ops.append([start, pos, cur])
    text = "".join(t for _k, t, _o in out)
    trimmed = []
    for lo, hi, k in ops:                 # don't let a range own trailing space
        while hi > lo and text[hi - 1].isspace():
            hi -= 1
        while lo < hi and text[lo].isspace():
            lo += 1
        if hi > lo:
            trimmed.append([lo, hi, k])
    return [[k, t] for k, t, _o in out], trimmed


def _idatui_unknown_row(ea, size):
    """One collapsed row for a run of ``size`` undefined bytes starting at
    ``ea``. A single byte is rendered normally (shows its value); a longer run
    collapses to ``db N dup(?)`` so a big .bss/gap doesn't explode into millions
    of one-byte rows."""
    import ida_name

    if size <= 1:
        return _idatui_head_row(ea)
    row = {"ea": hex(ea), "kind": "unknown", "size": int(size),
           "text": f"db {size} dup(?)"}
    nm = ida_name.get_ea_name(ea)
    if nm:
        row["name"] = nm
    return row


def _idatui_struct_member_rows(ea):
    """Indented member rows for a struct-typed data item at ``ea`` (expansion),
    or [] if it isn't a struct. Top-level fields only."""
    import ida_nalt
    import ida_typeinf
    import idaapi

    tif = ida_typeinf.tinfo_t()
    if not (ida_nalt.get_tinfo(tif, ea) and tif.is_udt()):
        return []
    udt = ida_typeinf.udt_type_data_t()
    if not tif.get_udt_details(udt):
        return []
    rows = []
    for m in udt:
        off = m.begin() // 8
        try:
            mtype = m.type._print() or ""
        except Exception:
            mtype = ""
        try:
            sz = int(m.type.get_size())
            if sz == idaapi.BADSIZE:
                sz = 0
        except Exception:
            sz = 0
        name = m.name or ""
        text = f"+{off:X} {name}" + (f" {mtype}" if mtype else "")
        rows.append({"ea": hex(ea + off), "kind": "member", "size": sz,
                     "text": text})
    return rows


def _idatui_func_header_rows(ea):
    """IDA-style subroutine banner rows shown just before a function's entry."""
    import ida_funcs

    name = ida_funcs.get_func_name(ea) or "sub_%X" % ea
    bar = "=" * 15 + " S U B R O U T I N E " + "=" * 15
    return [
        {"ea": hex(ea), "kind": "sep", "size": 0, "text": ""},
        {"ea": hex(ea), "kind": "sep", "size": 0, "text": "; " + bar},
        {"ea": hex(ea), "kind": "funchdr", "size": 0,
         "text": name + " proc", "name": name},
    ]


def _idatui_func_footer_rows(ea, func):
    """End-of-function marker shown just after a function's last item."""
    import ida_funcs

    name = ida_funcs.get_func_name(func.start_ea) or "sub_%X" % func.start_ea
    return [
        {"ea": hex(ea), "kind": "funchdr", "size": 0,
         "text": name + " endp", "name": name},
        {"ea": hex(ea), "kind": "sep", "size": 0, "text": "; " + "-" * 60},
    ]


@tool
@idasync
def heads(
    addr: Annotated[str, "Start address or name to walk from"],
    count: Annotated[int, "Max heads to return (default 200, max 2000)"] = 200,
    offset: Annotated[int, "Skip first N heads from addr (default 0)"] = 0,
    end: Annotated[str, "Optional exclusive end address; default = segment end"] = "",
    back: Annotated[bool, "Walk backwards: return the count heads ENDING just before addr, in forward order"] = False,
    annotate: Annotated[bool, "Emit IDA-style function boundary banner rows (kind sep/funchdr)"] = False,
) -> dict:
    """Walk item heads from ``addr`` as a flat listing: every head is rendered
    (code OR data OR undefined) via generate_disasm_line and stepped with
    next_head/prev_head. Unlike ``disasm`` (code-only, bails at the first data
    byte) this shows db/dw/dd/... lines for data and undefined regions — IDA's
    real disassembly view. Address-paged: page forward by re-calling with
    ``addr`` = the returned cursor.next; page up with ``back=true``."""
    import ida_bytes
    import ida_segment
    import idaapi

    count = 2000 if count > 2000 else (1 if count < 1 else count)
    offset = max(int(offset), 0)
    try:
        start = parse_address(addr)
    except Exception as e:
        return {"addr": str(addr), "error": str(e), "heads": [], "cursor": {"done": True}}
    seg = ida_segment.getseg(start)
    if not seg:
        return {"addr": str(addr), "error": "no segment", "heads": [], "cursor": {"done": True}}
    lo, hi = seg.start_ea, seg.end_ea
    if end:
        try:
            hi = min(hi, parse_address(end))
        except Exception:
            pass

    rows = []
    if back:
        # Collect up to (count+offset) heads strictly before `start`, then take
        # the window closest to `start`, returned in forward order.
        walk = []
        cur = ida_bytes.prev_head(start, lo)
        while cur != idaapi.BADADDR and cur >= lo and len(walk) < count + offset:
            walk.append(cur)
            cur = ida_bytes.prev_head(cur, lo)
        walk.reverse()
        chosen = walk[: len(walk) - offset] if offset else walk
        chosen = chosen[-count:]
        rows = [_idatui_head_row(e) for e in chosen]
        first = chosen[0] if chosen else start
        pea = ida_bytes.prev_head(first, lo)
        cursor = {"done": True} if pea == idaapi.BADADDR or pea < lo else {"prev": hex(pea)}
        return {"addr": str(addr), "heads": rows, "cursor": cursor}

    # Walk by item END (not next_head): next_head SKIPS undefined bytes, but a
    # flat listing must show them (IDA renders undefined as `db ?` lines, and
    # navigating to an unmarked address must land ON it). Defined items advance
    # by get_item_end; a run of undefined bytes is COLLAPSED into one row (its
    # end found in O(1) via next_head, which skips undefined) so a large .bss or
    # gap doesn't explode into millions of one-byte rows.
    def _is_unknown(e):
        f = ida_bytes.get_flags(e)
        return not (ida_bytes.is_code(f) or ida_bytes.is_data(f))

    def _run_end(e):
        """End (exclusive) of the undefined run starting at ``e``."""
        nh = ida_bytes.next_head(e, hi)
        return nh if (nh != idaapi.BADADDR and e < nh <= hi) else hi

    def _advance(e):
        if _is_unknown(e):
            return _run_end(e)
        nxt = ida_bytes.get_item_end(e)
        return nxt if nxt > e else e + 1

    def _rows_for(e):
        if _is_unknown(e):
            return [_idatui_unknown_row(e, _run_end(e) - e)]
        func = idaapi.get_func(e) if annotate else None
        at_start = func is not None and func.start_ea == e
        out = []
        if at_start:
            out.extend(_idatui_func_header_rows(e))
        row = _idatui_head_row(e)
        if at_start:
            row = dict(row)
            row["name"] = None  # the name is shown on the proc header line
        elif annotate and row.get("kind") == "code" and row.get("name"):
            # A code label (loc_XXX/jump target) gets its OWN line at depth 0,
            # like IDA; strip it from the instruction row below.
            nm = row["name"]
            out.append({"ea": hex(e), "kind": "label", "size": 0,
                        "text": nm + ":", "name": nm})
            row = dict(row)
            row["name"] = None
        out.append(row)
        if row.get("kind") == "data":
            out.extend(_idatui_struct_member_rows(e))  # expand struct fields
        if func is not None and ida_bytes.get_item_end(e) >= func.end_ea:
            out.extend(_idatui_func_footer_rows(e, func))
        return out

    ea = ida_bytes.get_item_head(start)
    for _ in range(offset):
        if ea >= hi or ea == idaapi.BADADDR:
            break
        ea = _advance(ea)
    more = False
    while ea != idaapi.BADADDR and ea < hi:
        if len(rows) >= count:
            more = True
            break
        rows.extend(_rows_for(ea))  # a struct head expands into member rows
        ea = _advance(ea)
    cursor = {"next": hex(ea)} if more else {"done": True}
    return {"addr": str(addr), "heads": rows, "cursor": cursor}


@tool
@idasync
def xref_types(
    queries: Annotated[list, "[{addr, direction:'to'|'from'|'both', include_fn, dedup, count}]"],
) -> dict:
    """Like xref_query, but every row carries a fine-grained ``kind`` derived from
    the IDA xref type \u2014 call/jump/flow for code, read/write/offset/text/info for
    data \u2014 alongside the coarse ``type`` (code/data). Feeds the xref dialog's
    r/w/call badges. Same query/envelope shape as xref_query."""
    import idaapi, idautils, ida_funcs, ida_bytes, ida_xref
    code_kind = {ida_xref.fl_CF: "call", ida_xref.fl_CN: "call",
                 ida_xref.fl_JF: "jump", ida_xref.fl_JN: "jump",
                 ida_xref.fl_F: "flow"}
    data_kind = {ida_xref.dr_O: "offset", ida_xref.dr_W: "write",
                 ida_xref.dr_R: "read", ida_xref.dr_T: "text", ida_xref.dr_I: "info"}

    def _kind(xr):
        table = code_kind if xr.iscode else data_kind
        return table.get(xr.type, "code" if xr.iscode else "data")

    def _fn(ea):
        f = ida_funcs.get_func(ea)
        if not f:
            return None
        return {"addr": hex(f.start_ea), "name": ida_funcs.get_func_name(f.start_ea)}

    def _resolve(raw):
        raw = str(raw).strip()
        try:
            return int(raw, 16)  # handles '0x2490' and '2490'
        except ValueError:
            return idaapi.get_name_ea(idaapi.BADADDR, raw)

    qs = queries if isinstance(queries, list) else [queries]
    result = []
    for q in qs:
        q = q if isinstance(q, dict) else {"addr": q}
        raw = str(q.get("addr", "")).strip()
        direction = str(q.get("direction", "to") or "to").lower()
        include_fn = bool(q.get("include_fn", True))
        dedup = bool(q.get("dedup", True))
        try:
            count = int(q.get("count", 2000) or 2000)
        except (TypeError, ValueError):
            count = 2000
        target = _resolve(raw)
        rows = []
        if target is not None and target != idaapi.BADADDR and ida_bytes.is_mapped(target):
            if direction in ("to", "both"):
                for xr in idautils.XrefsTo(target, 0):
                    row = {"direction": "to", "addr": hex(xr.frm), "from": hex(xr.frm),
                           "to": hex(target), "type": "code" if xr.iscode else "data",
                           "kind": _kind(xr)}
                    if include_fn:
                        row["fn"] = _fn(xr.frm)
                    rows.append(row)
            if direction in ("from", "both"):
                for xr in idautils.XrefsFrom(target, 0):
                    row = {"direction": "from", "addr": hex(xr.to), "from": hex(target),
                           "to": hex(xr.to), "type": "code" if xr.iscode else "data",
                           "kind": _kind(xr)}
                    if include_fn:
                        row["fn"] = _fn(xr.to)
                    rows.append(row)
            if dedup:
                seen = set()
                deduped = []
                for r in rows:
                    k = (r["direction"], r["from"], r["to"], r["kind"])
                    if k in seen:
                        continue
                    seen.add(k)
                    deduped.append(r)
                rows = deduped
            rows = rows[:count]
        result.append({"query": raw, "data": rows, "next_offset": None})
    return {"result": result}


@tool
@idasync
def data_type(
    addr: Annotated[str, "Address or name of a data item / global"],
) -> dict:
    """The current C type of a data item, for prefilling a retype prompt:
    {addr, name, type, size, is_func}. ``type`` is empty when the item is
    untyped; ``is_func`` distinguishes a global from a function so the caller
    knows which flavour of set_type to use."""
    import idaapi
    import ida_bytes
    import ida_name
    import idc
    raw = str(addr).strip()
    try:
        ea = int(raw, 16)
    except ValueError:
        ea = idaapi.get_name_ea(idaapi.BADADDR, raw)
    if ea == idaapi.BADADDR or not ida_bytes.is_mapped(ea):
        return {"addr": raw, "error": f"not a mapped address: {raw}"}
    return {
        "addr": hex(ea),
        "name": ida_name.get_name(ea) or "",
        "type": idc.get_type(ea) or "",
        "size": int(ida_bytes.get_item_size(ea) or 0),
        "is_func": bool(idaapi.get_func(ea)),
    }


@tool
@idasync
def decomp_map(
    addr: Annotated[str, "Function address or name"],
) -> dict:
    """Per-pseudocode-line instruction coverage for the split view's region
    highlight: for each line, the set of EAs the decompiler attributes to it,
    swept across the line's columns via get_line_item. Shape:
    {addr, lines:[{ea: primary|None, eas:[hex,...]}, ...]}."""
    import ida_hexrays
    import idaapi
    try:
        ea = int(str(addr), 16)
    except ValueError:
        ea = idaapi.get_name_ea(idaapi.BADADDR, str(addr).strip())
    func = idaapi.get_func(ea)
    if not func:
        return {"error": f"no function at {addr}"}
    try:
        cfunc = ida_hexrays.decompile(func.start_ea)
    except Exception as e:  # noqa: BLE001
        return {"error": f"decompile failed: {e}"}
    if cfunc is None:
        return {"error": "decompile failed"}
    lines = []
    for sl in cfunc.get_pseudocode():
        line = sl.line
        eas, seen = [], set()
        for x in range(len(line) + 1):
            head = ida_hexrays.ctree_item_t()
            item = ida_hexrays.ctree_item_t()
            tail = ida_hexrays.ctree_item_t()
            if not cfunc.get_line_item(line, x, False, head, item, tail):
                continue
            # Match the /*ea*/ marker's source (decompile_function_safe): the
            # item's dstr() is 'EA: description'; get_ea() reports a different ea.
            dstr = item.dstr()
            if not dstr:
                continue
            parts = dstr.split(": ", 1)
            if len(parts) != 2:
                continue
            try:
                e = int(parts[0], 16)
            except ValueError:
                continue
            if e not in seen:
                seen.add(e)
                eas.append(hex(e))
        lines.append({"ea": eas[0] if eas else None, "eas": eas})
    return {"addr": hex(func.start_ea), "lines": lines}


_idatui_strings_cache = {}


def _idatui_build_strings(min_len):
    """[(ea, text, length, typename)] for every string IDA found, cached by
    min_len (rebuilding the list is O(n) and the browser pages through it)."""
    import idautils
    import ida_nalt
    hit = _idatui_strings_cache.get(min_len)
    if hit is not None:
        return hit
    tnames = {}
    for nm, lbl in (("STRTYPE_C", "C"), ("STRTYPE_C_16", "utf16"),
                    ("STRTYPE_C_32", "utf32"), ("STRTYPE_PASCAL", "pascal")):
        v = getattr(ida_nalt, nm, None)
        if v is not None:
            tnames[v & 0xFF] = lbl
    items = []
    for s in idautils.Strings():
        if s is None:
            continue
        try:
            text = str(s)
        except Exception:  # noqa: BLE001 -- undecodable literal
            continue
        if len(text) < min_len:
            continue
        st = getattr(s, "strtype", 0) & 0xFF
        items.append((s.ea, text, getattr(s, "length", len(text)),
                      tnames.get(st, "t%d" % st)))
    _idatui_strings_cache[min_len] = items
    return items


@tool
@idasync
def list_strings(
    offset: Annotated[int, "Start index into the strings list"] = 0,
    count: Annotated[int, "Max strings to return (page size)"] = 2000,
    min_len: Annotated[int, "Minimum string length to include"] = 4,
    refresh: Annotated[bool, "Rebuild the cached strings list"] = False,
) -> dict:
    """Every string literal IDA found in the binary (IDA's Shift+F12 window),
    paginated: {strings:[{addr,text,len,type}], total, next_offset}. Feeds the
    TUI's strings browser."""
    try:
        min_len = max(int(min_len), 1)
    except (TypeError, ValueError):
        min_len = 4
    try:
        offset = max(int(offset), 0)
    except (TypeError, ValueError):
        offset = 0
    try:
        count = max(int(count), 1)
    except (TypeError, ValueError):
        count = 2000
    if refresh:
        _idatui_strings_cache.pop(min_len, None)
    items = _idatui_build_strings(min_len)
    page = items[offset:offset + count]
    return {
        "strings": [{"addr": hex(ea), "text": text, "len": ln, "type": ty}
                    for (ea, text, ln, ty) in page],
        "total": len(items),
        "next_offset": offset + len(page),
    }

@tool
@idasync
def list_linkage(
    kind: Annotated[str, "'import', 'export' or 'both'"] = "both",
) -> dict:
    """What this binary imports from, and exports to, other modules:
    {imports:[{addr,name,module}], exports:[{addr,name,ordinal}]}. Feeds the
    project-wide import/export join, which resolves a PLT stub in one binary to
    the real implementation in another."""
    import idaapi
    import idautils
    import ida_nalt
    want = str(kind or "both").lower()
    imports = []
    exports = []
    if want in ("import", "both"):
        n = ida_nalt.get_import_module_qty()
        for i in range(n):
            mod = ida_nalt.get_import_module_name(i) or ""

            def _cb(ea, name, ordinal, _mod=mod):
                # An ordinal-only import has no name; skip rather than invent one.
                if name:
                    imports.append({"addr": hex(ea), "name": name, "module": _mod})
                return True

            ida_nalt.enum_import_names(i, _cb)
    if want in ("export", "both"):
        for index, ordinal, ea, name in idautils.Entries():
            if name:
                exports.append({"addr": hex(ea), "name": name,
                                "ordinal": int(ordinal)})
    return {"imports": imports, "exports": exports,
            "n_imports": len(imports), "n_exports": len(exports)}

@tool
@idasync
def define_code_run(
    addr: Annotated[str, "Address to start disassembling from"],
    limit: Annotated[int, "Max instructions to create (safety stop)"] = 20000,
) -> dict:
    """Disassemble CONSECUTIVELY from ``addr`` until something stops it, the way
    IDA's 'c' does — one instruction is rarely what you want when carving a raw
    image. Returns {start,end,count,stopped} where ``stopped`` says why:
    'undecodable' (bytes aren't an instruction), 'flow' (the last instruction
    doesn't fall through, e.g. RET/B), 'defined' (ran into existing code/data),
    'segment' (hit the end) or 'limit'.

    Runs in-process: doing this from the client would be one round trip per
    instruction, which is minutes on a real firmware image."""
    import ida_bytes
    import ida_idp
    import ida_segment
    import ida_ua
    import idaapi

    try:
        ea = parse_address(addr)
    except Exception as e:
        return {"addr": str(addr), "error": str(e), "count": 0}

    seg = ida_segment.getseg(ea)
    if not seg:
        return {"addr": str(addr), "error": "no segment", "count": 0}
    hi = seg.end_ea
    try:
        limit = max(1, min(int(limit), 200000))
    except (TypeError, ValueError):
        limit = 20000

    start, count, stopped = ea, 0, "limit"
    while count < limit:
        if ea >= hi:
            stopped = "segment"
            break
        flags = ida_bytes.get_flags(ea)
        if ida_bytes.is_code(flags) or ida_bytes.is_data(flags):
            # Already defined: stop rather than clobber. Undefining someone's
            # existing work to keep a speculative run going is not a trade the
            # user asked for.
            stopped = "defined"
            break
        n = ida_ua.create_insn(ea)
        if n <= 0:
            stopped = "undecodable"
            break
        count += 1
        # Stop where control flow stops. Past a RET the next bytes are usually
        # padding or a new function's data, and running on turns a clean carve
        # into a mess that has to be undone by hand.
        #
        # Ask ida_idp.is_ret_insn, NOT the canonical feature bits: on AArch64
        # get_canon_feature() returns 0 for RET, so a CF_STOP test silently never
        # fires and the run walks straight through the end of the routine.
        insn = ida_ua.insn_t()
        if ida_ua.decode_insn(insn, ea) > 0:
            try:
                is_ret = ida_idp.is_ret_insn(insn)
            except Exception:
                is_ret = False
            if is_ret or (insn.get_canon_feature() & idaapi.CF_STOP):
                ea += n
                stopped = "flow"
                break
        ea += n

    return {"start": hex(start), "end": hex(ea), "count": count,
            "stopped": stopped}

@tool
@idasync
def set_thumb(
    addr: Annotated[str, "Address to change the ARM decoding mode at"],
    mode: Annotated[str, "'toggle', 'on' (Thumb) or 'off' (ARM)"] = "toggle",
    end: Annotated[str, "Optional exclusive end address (default: this item)"] = "",
) -> dict:
    """Switch ARM/Thumb decoding at ``addr`` (IDA's T segment register).

    Thumb is not a property of the bytes, it's a mode the CPU is in, so a raw
    image gives IDA no way to know: at a Thumb entry point it decodes 16-bit
    instructions as 32-bit ARM and produces confident nonsense
    (``push {r3,lr}`` reads as ``SVCLT 0xBF00``).

    Also forces the segment to 32-bit when turning Thumb ON. Thumb does not
    exist in AArch64, and a headerless blob loaded with -parm defaults to
    64-bit — so setting T alone changes nothing and looks broken. Asking for
    Thumb IS asking for ARM32."""
    import ida_bytes
    import ida_idp
    import ida_segment
    import ida_segregs

    try:
        ea = parse_address(addr)
    except Exception as e:
        return {"addr": str(addr), "error": str(e)}
    treg = ida_idp.str2reg("T")
    if treg is None or treg < 0:
        return {"addr": hex(ea), "error": "no T register (not an ARM database)"}
    seg = ida_segment.getseg(ea)
    if not seg:
        return {"addr": hex(ea), "error": "no segment"}

    import ida_ida
    db64 = ida_ida.inf_get_app_bitness() == 64
    cur = ida_segregs.get_sreg(ea, treg)
    cur = 0 if cur in (None, 0xFFFFFFFF, -1) else int(cur)
    want = {"on": 1, "off": 0}.get(str(mode).lower(), 0 if cur else 1)

    changed_bits = False
    if want and seg.bitness != 1:
        ida_segment.set_segm_addressing(seg, 1)
        changed_bits = True

    try:
        stop = parse_address(end) if end else 0
    except Exception:
        stop = 0
    size = max(int(stop) - ea, 0) or max(ida_bytes.get_item_size(ea), 2)
    # The bytes are currently decoded in the OLD mode; leaving that item defined
    # pins the wrong instruction length and the new mode has nothing to apply to.
    ida_bytes.del_items(ea, 0, size)
    ok = bool(ida_segregs.split_sreg_range(ea, treg, want, ida_segregs.SR_user))
    now = ida_segregs.get_sreg(ea, treg)
    return {"addr": hex(ea), "thumb": bool(now), "was": bool(cur), "ok": ok,
            "bitness": ida_segment.getseg(ea).bitness,
            "forced_32bit": changed_bits,
            # The DATABASE's bitness is fixed at load and can't be corrected
            # here (setting it post-hoc makes the decompiler INTERR). In a
            # 64-bit database a 32-bit function disassembles but Hex-Rays
            # refuses it outright, so say so instead of leaving the user to
            # discover that F5 does nothing.
            "db_64bit": bool(db64 and want)}

def _idatui_add_func(ea):
    """add_func at ``ea``, falling back to an explicit end.

    ida_funcs.add_func(ea) asks IDA to find the end and on carved or
    freshly-marked code it often can't, failing with no reason given."""
    import ida_bytes
    import ida_funcs
    import ida_segment
    import idaapi

    if idaapi.get_func(ea) is not None:
        return True
    if ida_funcs.add_func(ea):
        return True
    seg = ida_segment.getseg(ea)
    hi = seg.end_ea if seg else ea
    end = ea
    while end < hi and ida_bytes.is_code(ida_bytes.get_flags(end)):
        nxt = ida_bytes.get_item_end(end)
        if nxt <= end:
            break
        end = nxt
    return bool(end > ea and ida_funcs.add_func(ea, end))


@tool
@idasync
def define_func_run(
    addr: Annotated[str, "Entry point of the function to create"],
) -> dict:
    """Create a function at ``addr``, working out its end if IDA can't.

    ida_funcs.add_func(ea) asks IDA to find the end itself, and on hand-carved
    code it often can't — a run that ends in a tail call, or whose last
    instruction isn't recognised as a return, simply fails with no reason given.
    You then have a disassembled routine that refuses to become a function, and
    F5 has nothing to work with.

    So: try IDA's way, and if that fails, use the end of the contiguous
    instruction run starting at ``addr``."""
    import ida_bytes
    import ida_funcs
    import ida_segment
    import idaapi

    try:
        ea = parse_address(addr)
    except Exception as e:
        return {"addr": str(addr), "error": str(e), "ok": False}
    fn = idaapi.get_func(ea)
    if fn is not None and fn.start_ea == ea:
        return {"addr": hex(ea), "ok": True, "start": hex(fn.start_ea),
                "end": hex(fn.end_ea), "how": "existed"}
    auto = ida_funcs.add_func(ea)
    if not auto and not _idatui_add_func(ea):
        return {"addr": hex(ea), "ok": False,
                "error": f"IDA refused a function at {ea:#x}"}
    f = idaapi.get_func(ea)
    if f is None:
        return {"addr": hex(ea), "ok": False, "error": "function did not stick"}
    return {"addr": hex(ea), "ok": True, "start": hex(f.start_ea),
            "end": hex(f.end_ea), "how": "auto" if auto else "explicit-end"}

@tool
@idasync
def decomp_error(
    addr: Annotated[str, "Address of the function that failed to decompile"],
) -> dict:
    """Why Hex-Rays refused this function, in its own words.

    The plain decompile tool reports "Decompilation failed at 0x0" and drops the
    reason, which is the only useful part. Hex-Rays fills in a hexrays_failure_t
    saying things like "only 64-bit functions can be decompiled in the current
    database" — that one is unfixable in place (the database's bitness is set at
    load), so a user who can't see it has no way to know they must reload."""
    import ida_funcs
    import ida_hexrays
    import ida_ida

    try:
        ea = parse_address(addr)
    except Exception as e:
        return {"addr": str(addr), "error": str(e)}
    out = {"addr": hex(ea), "bitness": ida_ida.inf_get_app_bitness()}
    fn = ida_funcs.get_func(ea)
    if fn is None:
        out["reason"] = "no function here"
        return out
    try:
        if not ida_hexrays.init_hexrays_plugin():
            out["reason"] = "the decompiler is not available for this processor"
            return out
        hf = ida_hexrays.hexrays_failure_t()
        cf = ida_hexrays.decompile_func(fn, hf)
        if cf is not None:
            out["reason"] = ""      # it decompiles now
            return out
        out["reason"] = hf.desc() or f"error {hf.code}"
        out["code"] = int(hf.code)
        out["errea"] = hex(hf.errea)
    except Exception as e:  # noqa: BLE001
        out["reason"] = f"{type(e).__name__}: {e}"
    return out

@tool
@idasync
def thumb_scan(
    start: Annotated[str, "Start of the range to scan for entry pointers"] = "",
    end: Annotated[str, "Exclusive end of the range (default: 1KB from start)"] = "",
    apply: Annotated[bool, "Mark the targets as Thumb and disassemble them"] = True,
    limit: Annotated[int, "Max entries to act on"] = 512,
) -> dict:
    """Find Thumb entry points from ODD pointers, e.g. a Cortex-M vector table.

    An ARM function pointer carries the mode in bit 0: odd means Thumb. A vector
    table is therefore a list of Thumb entry points that IDA won't follow on a
    headerless image, because nothing tells it those words are pointers at all.

    Being wrong here is expensive — marking a data word as code corrupts the
    listing — so a word only counts when it is odd, lands inside a loaded
    segment, and its target is EXECUTABLE and not already defined as data. The
    even words in a vector table (the initial stack pointer) fail the first test,
    which is the point."""
    import ida_bytes
    import ida_funcs
    import ida_idp
    import ida_segment
    import ida_segregs
    import ida_ua

    seg0 = ida_segment.getseg(parse_address(start)) if start else None
    if seg0 is None:
        seg0 = ida_segment.getnseg(0)
    if seg0 is None:
        return {"error": "no segments", "found": [], "applied": 0}
    try:
        lo = parse_address(start) if start else seg0.start_ea
        hi = parse_address(end) if end else min(lo + 0x400, seg0.end_ea)
    except Exception as e:
        return {"error": str(e), "found": [], "applied": 0}

    treg = ida_idp.str2reg("T")
    found, applied = [], 0
    ea = lo
    while ea + 4 <= hi and len(found) < limit:
        w = ida_bytes.get_dword(ea)
        ea += 4
        if not (w & 1):
            continue                       # even: not a Thumb pointer
        tgt = w & ~1
        seg = ida_segment.getseg(tgt)
        if seg is None or not (seg.perm & ida_segment.SEGPERM_EXEC or seg.perm == 0):
            continue                       # points outside the image, or at data
        f = ida_bytes.get_flags(tgt)
        if ida_bytes.is_data(f):
            continue                       # already something else; don't fight it
        rec = {"at": hex(ea - 4), "value": hex(w), "target": hex(tgt),
               "was_code": bool(ida_bytes.is_code(f))}
        found.append(rec)
        if not apply:
            continue
        if treg is not None and treg >= 0:
            ida_segregs.split_sreg_range(tgt, treg, 1, ida_segregs.SR_user)
        if not ida_bytes.is_code(ida_bytes.get_flags(tgt)):
            ida_bytes.del_items(tgt, 0, 2)
            if ida_ua.create_insn(tgt) <= 0:
                rec["decoded"] = False
                continue
        rec["decoded"] = True
        rec["function"] = _idatui_add_func(tgt)
        applied += 1
    return {"start": hex(lo), "end": hex(hi), "found": found,
            "applied": applied, "n": len(found)}


# --------------------------------------------------------------------------- #
# operand display formats (IDA's 'o' family: hex / dec / char / offset / ...)
# --------------------------------------------------------------------------- #
#: The stops a cycle walks, in order, before filtering to the ones that make
#: sense for the operand in hand. Octal is deliberately NOT one of them -- every
#: extra stop is another keypress and nobody reads octal -- but it is still
#: reachable by name. "default" hands the operand back to IDA's own choice,
#: which for data is how you get an auto-detected offset/string back.
_IDATUI_FMT_CYCLE = ("hex", "dec", "bin", "char", "offset", "default")

#: Formats we can re-apply from a name alone. enum/stroff/custom carry an id
#: (which enum, which struct) that a nibble doesn't record, so they are never
#: cycled INTO -- and cycling out of one is called out in ``warn``.
_IDATUI_FMT_SETTABLE = ("hex", "dec", "oct", "bin", "char", "offset", "seg",
                        "float", "stack", "default")


def _idatui_fmt_nibbles():
    """{format name: IDA operand-type nibble}. Built on call, not at import:
    this module is injected into a file that is imported before a database is
    open."""
    import ida_bytes
    return {
        "default": ida_bytes.FF_N_VOID, "hex": ida_bytes.FF_N_NUMH,
        "dec": ida_bytes.FF_N_NUMD, "char": ida_bytes.FF_N_CHAR,
        "seg": ida_bytes.FF_N_SEG, "offset": ida_bytes.FF_N_OFF,
        "bin": ida_bytes.FF_N_NUMB, "oct": ida_bytes.FF_N_NUMO,
        "enum": ida_bytes.FF_N_ENUM, "forced": ida_bytes.FF_N_FOP,
        "stroff": ida_bytes.FF_N_STRO, "stack": ida_bytes.FF_N_STK,
        "float": ida_bytes.FF_N_FLT, "custom": ida_bytes.FF_N_CUST,
    }


def _idatui_fmt_name(nib):
    for name, v in _idatui_fmt_nibbles().items():
        if v == nib:
            return name
    return "default"


def _idatui_op_fmt(ea, n):
    """The format operand ``n`` of the item at ``ea`` is currently displayed in.

    Reads the nibble IDA keeps per operand rather than guessing from the text --
    ``1`` renders identically in hex and decimal, so the rendered line cannot
    answer this."""
    import ida_bytes
    F = ida_bytes.get_flags(ea)
    nib = (F >> ida_bytes.get_operand_type_shift(int(n))) & 0xF
    return _idatui_fmt_name(nib)


def _idatui_op_value(ea, n):
    """(value, byte width) of operand ``n``, or (None, 0) if it hasn't got one.

    The value is what decides which formats are OFFERED: a character constant
    for 0x38A9 or an offset to an unmapped address are stops worth skipping."""
    import ida_bytes
    import ida_ua

    F = ida_bytes.get_flags(ea)
    if ida_bytes.is_code(F):
        insn = ida_ua.insn_t()
        if ida_ua.decode_insn(insn, ea) <= 0:
            return None, 0
        try:
            op = insn.ops[int(n)]
        except Exception:
            return None, 0
        if op.type == ida_ua.o_void:
            return None, 0
        v = op.value if op.type == ida_ua.o_imm else op.addr
        try:
            size = int(ida_ua.get_dtype_size(op.dtype))
        except Exception:
            size = 0
        return int(v), size
    size = int(ida_bytes.get_item_size(ea))
    read = {1: ida_bytes.get_byte, 2: ida_bytes.get_word,
            4: ida_bytes.get_dword, 8: ida_bytes.get_qword}.get(size)
    if read is None:
        return None, size
    try:
        return int(read(ea)), size
    except Exception:
        return None, size


def _idatui_printable(v):
    """Whether ``v`` would actually render as a character constant. IDA accepts
    op_chr on anything and then prints the number anyway, so a cycle that offers
    'char' for 0x18 has a stop where nothing visibly happens."""
    if v is None or v < 0 or v > 0xFFFFFFFF:
        return False
    bs, x = [], int(v)
    while True:
        bs.append(x & 0xFF)
        x >>= 8
        if not x:
            break
    return all(0x20 <= b <= 0x7E or b in (9, 10, 13) for b in bs)


def _idatui_offset_worth(v):
    """Whether 'offset' is worth OFFERING as a cycle stop for value ``v``.

    Making an offset is not free: IDA invents a dummy name at the target
    (``off_18``) and that name STAYS once you cycle past it. So the ring only
    stops there when the target is already something you could name -- a symbol,
    a function, or an item something else references. In a PIE at base 0 half
    the small constants in a function are 'mapped' (they land in the ELF
    header); ``sub rsp, 18h`` is not a reference and must not offer to become
    one on the way past.

    An explicit request still converts anything mapped: that's a decision, not a
    keypress that happened to land here. After it, the target HAS a name, so the
    ring includes the stop from then on."""
    import ida_bytes
    import ida_name

    return bool(v and ida_bytes.is_mapped(v) and ida_name.get_ea_name(v))


def _idatui_op_candidates(ea):
    """Operand indices at ``ea`` whose display format is worth changing.

    Immediates and displacements -- the literals. Deliberately NOT:

    * branch targets (o_near/o_far), or every jump on the listing would offer to
      become a bare number, on a view you navigate by label;
    * memory references (o_mem), e.g. x86-64's RIP-relative ``lea rdi, name``.
      IDA prints those from the reference, not from the operand's number format,
      so setting one is accepted and changes nothing on screen -- a keypress
      that appears to do nothing is worse than one that says it can't.

    An explicit ``n`` still reaches them; this is what a bare cursor picks."""
    import ida_bytes
    import ida_ua

    F = ida_bytes.get_flags(ea)
    if ida_bytes.is_data(F):
        return [0]              # a data item's value is operand 0
    if not ida_bytes.is_code(F):
        return []               # undefined bytes: IDA refuses a format outright
    insn = ida_ua.insn_t()
    if ida_ua.decode_insn(insn, ea) <= 0:
        return []
    want = (ida_ua.o_imm, ida_ua.o_displ)
    out = []
    for i in range(len(insn.ops)):
        op = insn.ops[i]
        if op.type == ida_ua.o_void:
            break
        if op.type in want:
            out.append(i)
    return out


def _idatui_op_spans(ea, text):
    """[(start, end, n)] -- where each operand sits inside ``text`` (the
    whitespace-collapsed line the TUI shows), so a cursor column can name the
    operand it is standing on.

    Read out of IDA's own COLOR_OPND markers on the line, which is both free
    (the line is generated anyway) and exact. print_operand is kept as a
    fallback for a processor module that emits no operand markers -- it agrees
    with the tags where both exist, but it re-renders every operand to say so.
    """
    import ida_lines
    import ida_ua

    line = ida_lines.generate_disasm_line(ea, 0)
    if line:
        _spans, ops = _idatui_spans(line)
        if ops:
            return [tuple(o) for o in ops]

    out, pos = [], 0
    for n in range(8):
        try:
            raw = ida_ua.print_operand(ea, n)
        except Exception:
            raw = None
        if not raw:
            continue
        op = " ".join(ida_lines.tag_remove(raw).split())
        if not op:
            continue
        i = text.find(op, pos)
        if i < 0:                      # duplicated operand text (mov eax, eax)
            i = text.find(op)
        if i < 0:
            continue
        out.append((i, i + len(op), n))
        pos = i + len(op)
    return out


def _idatui_line_text(ea):
    import ida_lines
    line = ida_lines.generate_disasm_line(ea, 0)
    return " ".join(ida_lines.tag_remove(line).split()) if line else ""


def _idatui_op_text(ea, text, n):
    """How operand ``n`` reads on the line, for a message that names it."""
    for lo, hi, i in _idatui_op_spans(ea, text):
        if i == int(n):
            return text[lo:hi].strip()
    return ""


def _idatui_apply_fmt(ea, n, fmt):
    """Set operand ``n``'s display format. Returns (ok, error)."""
    import ida_bytes
    import ida_offset
    import idaapi

    n = int(n)
    if fmt == "default":
        return bool(ida_bytes.clr_op_type(ea, n)), ""
    if fmt == "offset":
        base = ida_offset.calc_offset_base(ea, n)
        if base in (idaapi.BADADDR, None) or base < 0:
            base = 0
        return bool(ida_offset.op_plain_offset(ea, n, base)), ""
    fn = {"hex": ida_bytes.op_hex, "dec": ida_bytes.op_dec,
          "oct": ida_bytes.op_oct, "bin": ida_bytes.op_bin,
          "char": ida_bytes.op_chr, "seg": ida_bytes.op_seg,
          "float": ida_bytes.op_flt, "stack": ida_bytes.op_stkvar}.get(fmt)
    if fn is None:
        return False, (f"can't set {fmt!r} from a name alone"
                       if fmt in _idatui_fmt_nibbles() else
                       f"unknown format {fmt!r}")
    return bool(fn(ea, n)), ""


@tool
@idasync
def op_format(
    addr: Annotated[str, "Address of the instruction or data item"],
    mode: Annotated[str, "cycle | back | show | hex | dec | oct | bin | char | offset | stack | default"] = "cycle",
    col: Annotated[int, "Cursor column inside the rendered line (-1: first literal)"] = -1,
    n: Annotated[int, "Operand index; -1 derives it from ``col``"] = -1,
) -> dict:
    """Change how a literal is DISPLAYED (IDA's 'o' family): hex, decimal,
    binary, character, or an offset to the address it names.

    The value in the bytes never changes -- only the representation IDA renders
    and remembers. ``cycle``/``back`` step the stops that make sense for THIS
    operand: 'char' is skipped unless the value prints as one, 'offset' unless
    the target is already named, so no press is ever a no-op you have to press
    again. ``show`` reports without changing anything.

    A format the ring can't hold (a stack variable, an enum) is reported in
    ``warn`` on the way out, with what to do about it -- ``mode`` takes any of
    the names above outright, which is also how you put one back.

    Which operand: ``n`` if given, else the one under ``col`` (a column in the
    whitespace-collapsed line, as ``heads`` renders it), else the first literal
    on the line."""
    import ida_bytes

    try:
        ea = ida_bytes.get_item_head(parse_address(addr))
    except Exception as e:
        return {"addr": str(addr), "error": str(e)}

    before = _idatui_line_text(ea)
    cands = _idatui_op_candidates(ea)
    n = int(n)
    if n < 0:
        n = -1
        if int(col) >= 0:
            for lo, hi, i in _idatui_op_spans(ea, before):
                if not (lo <= int(col) < hi):
                    continue
                if i in cands:
                    n = i
                    break
                # The cursor IS on an operand, just not one with a format. The
                # client highlights what the cursor is on, so quietly moving to
                # a different operand would make that highlight a lie -- say
                # which one can be changed instead.
                where = before[lo:hi].strip()
                alt = (f"; the literal on this line is operand {cands[0]} "
                       f"({_idatui_op_text(ea, before, cands[0])})"
                       if cands else "")
                return {"addr": hex(ea), "n": i, "text": before,
                        "error": f"operand {i} ({where}) has no format to "
                                 f"change{alt}"}
        if n < 0:
            if not cands:
                F = ida_bytes.get_flags(ea)
                why = ("no literal on this line to reformat"
                       if ida_bytes.is_code(F) or ida_bytes.is_data(F) else
                       "undefined bytes have no format to change -- define "
                       "them first ('d' makes data, 'c' makes code)")
                return {"addr": hex(ea), "text": before, "error": why}
            n = cands[0]

    cur = _idatui_op_fmt(ea, n)
    value, width = _idatui_op_value(ea, n)
    mapped = value is not None and value != 0 and ida_bytes.is_mapped(value)
    # The ring is a property of the OPERAND, not of what you last pressed: every
    # stop is one that changes what you see for this value, and it is the same
    # ring at every step, so a lap always comes home.
    choices = [f for f in _IDATUI_FMT_CYCLE
               if (f != "char" or _idatui_printable(value))
               and (f != "offset" or _idatui_offset_worth(value))]
    # A stack variable is deliberately NOT a stop: ``[rbp+var_40]`` is a frame
    # member, not a way of writing a number, and IDA's own "is this a stack
    # variable" test isn't exposed to Python here (calc_stkvar_struc_offset
    # happily answers for ``[r14+8]`` too, which would put a bogus stop in the
    # ring). Leaving one is reported instead, with the command that undoes it.
    lossy = cur not in choices and cur != "default"

    mode = str(mode or "cycle").lower()
    if mode == "show":
        return {"addr": hex(ea), "n": n, "format": cur, "prev": cur,
                "choices": choices, "text": before, "before": before,
                "value": None if value is None else hex(value),
                "width": width, "applied": False}
    if mode in ("cycle", "back"):
        step = 1 if mode == "cycle" else -1
        if cur in choices:
            want = choices[(choices.index(cur) + step) % len(choices)]
        else:
            # Standing on a format the ring can't hold (an enum names a type a
            # nibble doesn't record): enter the ring at its end, don't skip a
            # stop working out where we "would have" been.
            want = choices[0] if step > 0 else choices[-1]
    else:
        want = mode
        if want not in _idatui_fmt_nibbles():
            return {"addr": hex(ea), "n": n, "text": before,
                    "error": f"unknown format {mode!r}; one of "
                             + ", ".join(_IDATUI_FMT_SETTABLE)}
        if want == "offset" and not mapped:
            return {"addr": hex(ea), "n": n, "text": before, "format": cur,
                    "error": (f"{'0x%x' % value if value is not None else 'this operand'}"
                              " isn't a mapped address -- an offset to it would"
                              " invent a name for nothing")}

    ok, err = _idatui_apply_fmt(ea, n, want)
    if err:
        return {"addr": hex(ea), "n": n, "text": before, "format": cur,
                "error": err}
    got = _idatui_op_fmt(ea, n)
    out = {"addr": hex(ea), "n": n, "prev": cur, "format": got,
           "requested": want, "applied": bool(ok), "choices": choices,
           "before": before, "text": _idatui_line_text(ea),
           "value": None if value is None else hex(value), "width": width}
    if not ok:
        out["error"] = f"IDA refused {want} on operand {n}"
    elif lossy:
        out["warn"] = (
            f"operand {n} was {cur} and the ring has no stop there -- "
            + (f"'{cur}' sets it again" if cur in _IDATUI_FMT_SETTABLE else
               f"{cur} names a type this can't put back, reassign it by hand"))
    return out


# --------------------------------------------------------------------------- #
# the same thing in the decompiler (Hex-Rays keeps its own number formats)
# --------------------------------------------------------------------------- #
#: Hex-Rays prints C, so two of the listing's stops are missing here: binary
#: (C has no binary literal -- the format takes and then renders decimal, which
#: would be a lie on screen) and offset (it makes the function fail to
#: decompile outright).
_IDATUI_PC_FMT_CYCLE = ("hex", "dec", "oct", "char", "default")


def _idatui_compact(line):
    """The ida-pro-mcp whitespace collapse the pseudocode is served through, so
    a column in what the client SHOWS can be mapped back to Hex-Rays' line."""
    try:
        from ida_pro_mcp.ida_mcp.utils import compact_whitespace
        return compact_whitespace(line)
    except Exception:
        import re as _re
        stripped = line.lstrip(" \t")
        lead = line[: len(line) - len(stripped)]
        return lead + _re.sub(r"[ \t]{2,}", " ", stripped)


def _idatui_compact_col(plain, compact, col):
    """The inverse of ``_idatui_uncompact_col``: a column in Hex-Rays' own line,
    expressed in the collapsed line the client shows."""
    j = 0
    for i in range(min(int(col), len(plain))):
        if j < len(compact) and plain[i] == compact[j]:
            j += 1
    return j


def _idatui_uncompact_col(plain, compact, col):
    """Map a column in the collapsed line back to the same character in the
    original. The transform only ever DELETES spaces, so walking both in step
    and skipping what vanished is exact."""
    i = 0
    for j in range(min(int(col), len(compact))):
        c = compact[j]
        while i < len(plain) and plain[i] != c:
            i += 1
        i += 1
    return min(i, max(len(plain) - 1, 0))


#: Characters that can be part of a C number literal as Hex-Rays prints one
#: (digits, hex letters, the 0x prefix, u/L suffixes).
_IDATUI_LIT_CHARS = frozenset("0123456789abcdefABCDEFxXuUlL")


def _idatui_lit_extent(plain, x):
    """The [start, end) of the literal token containing column ``x``.

    Hex-Rays says WHICH item a column belongs to, but not how wide the printed
    literal is -- and it attributes neighbouring punctuation to the same item,
    so ``if ( a1 > 1 )`` reports the closing paren as part of the number. The
    identity comes from the ctree; the extent is the run of literal characters
    around the column, which cannot reach a ``)`` or a space."""
    if x >= len(plain):
        return None
    if plain[x] == "'":                       # a character constant: '-'
        end = plain.find("'", x + 1)
        return (x, end + 1) if end > x else None
    lo = plain.rfind("'", 0, x)
    if lo >= 0 and plain.find("'", x) > x and "'" in plain[lo:x] and \
            plain[lo:x].count("'") == 1 and " " not in plain[lo:x]:
        return (lo, plain.find("'", x) + 1)   # inside 'c'
    if plain[x] not in _IDATUI_LIT_CHARS:
        return None
    lo = x
    while lo > 0 and plain[lo - 1] in _IDATUI_LIT_CHARS:
        lo -= 1
    hi = x
    while hi < len(plain) and plain[hi] in _IDATUI_LIT_CHARS:
        hi += 1
    if lo > 0 and plain[lo - 1] == "-":       # a unary minus is part of it
        lo -= 1
    return (lo, hi)


def _idatui_pc_nums(cf, sl):
    """Every number literal on one pseudocode line, as
    [{x0, x1, ea, opnum, value, nbytes, fmt}].

    Asks Hex-Rays what each column belongs to rather than pattern-matching the
    text: a regex over ``v6 = a1 - 1;`` has to guess which of those characters
    are a literal, and ``v11`` looks like one."""
    import ida_bytes
    import ida_hexrays
    import ida_lines
    import idaapi

    plain = ida_lines.tag_remove(sl.line)
    out = []
    x = 0
    while x < len(plain):
        ch = plain[x]
        if ch not in _IDATUI_LIT_CHARS and ch != "'":
            x += 1
            continue
        head, item, tail = (ida_hexrays.ctree_item_t() for _ in range(3))
        if not cf.get_line_item(sl.line, x, True, head, item, tail):
            x += 1
            continue
        if item.citype != ida_hexrays.VDI_EXPR:
            x += 1
            continue
        e = item.e
        if e.op != ida_hexrays.cot_num:
            x += 1
            continue
        extent = _idatui_lit_extent(plain, x)
        if extent is None:
            x += 1
            continue
        nf = e.n.nf
        opnum = ord(nf.opnum) if isinstance(nf.opnum, str) else int(nf.opnum)
        nbytes = (ord(nf.org_nbytes) if isinstance(nf.org_nbytes, str)
                  else int(nf.org_nbytes))
        ea = int(e.ea)
        if ea == idaapi.BADADDR:
            x = extent[1]
            continue                      # synthesised: nothing to key on
        nib = (nf.flags >> ida_bytes.get_operand_type_shift(opnum)) & 0xF
        # Whether this format is the USER's or Hex-Rays' own guess. The nibble
        # can't say: an untouched number reads back as whatever it happens to
        # be printed as, and cycling from there would skip that stop forever
        # (default already looks like it) and never come back to it.
        loc = ida_hexrays.operand_locator_t(ea, opnum)
        user = (ida_hexrays.user_numforms_find(cf.numforms, loc)
                != ida_hexrays.user_numforms_end(cf.numforms))
        out.append({"x0": extent[0], "x1": extent[1], "ea": ea,
                    "opnum": opnum, "value": int(e.n._value),
                    "nbytes": nbytes, "user": user,
                    "fmt": _idatui_fmt_name(nib) if user else "default",
                    "shown": _idatui_fmt_name(nib)})
        x = extent[1]                     # past this literal, not into it
    return out


@tool
@idasync
def pc_nums(
    addr: Annotated[str, "Function address (or any address inside it)"],
) -> dict:
    """Every number literal in a function's pseudocode, as
    [{line, x0, x1, ea, opnum, value, fmt, user}].

    One call per decompilation, so a client can show WHICH literal the cursor is
    on (and reformat exactly that one) without a round trip per cursor move.
    Columns are in the same collapsed coordinates the decompile tool serves its
    text in, i.e. what the client actually displays."""
    import ida_hexrays
    import ida_lines
    import idaapi

    if not ida_hexrays.init_hexrays_plugin():
        return {"addr": str(addr), "error": "no decompiler", "nums": []}
    try:
        f = idaapi.get_func(parse_address(addr))
    except Exception as e:
        return {"addr": str(addr), "error": str(e), "nums": []}
    if f is None:
        return {"addr": str(addr), "error": "no function here", "nums": []}
    try:
        cf = ida_hexrays.decompile(f.start_ea)
    except Exception as e:
        return {"addr": hex(f.start_ea), "error": f"decompile failed: {e}",
                "nums": []}
    if cf is None:
        return {"addr": hex(f.start_ea), "error": "decompilation failed",
                "nums": []}
    sv = cf.get_pseudocode()
    out = []
    for i in range(len(sv)):
        plain = ida_lines.tag_remove(sv[i].line)
        compact = _idatui_compact(plain)
        for rec in _idatui_pc_nums(cf, sv[i]):
            out.append({
                "line": i,
                "x0": _idatui_compact_col(plain, compact, rec["x0"]),
                "x1": _idatui_compact_col(plain, compact, rec["x1"]),
                "ea": hex(rec["ea"]), "opnum": rec["opnum"],
                "value": hex(rec["value"]), "fmt": rec["fmt"],
                "shown": rec["shown"], "user": bool(rec["user"]),
            })
    return {"addr": hex(f.start_ea), "nums": out, "lines": len(sv)}


@tool
@idasync
def pc_num_format(
    addr: Annotated[str, "Function address (or any address inside it)"],
    mode: Annotated[str, "cycle | back | show | hex | dec | oct | char | default"] = "cycle",
    line: Annotated[int, "0-based pseudocode line index"] = -1,
    col: Annotated[int, "Cursor column in the DISPLAYED line (-1: first literal)"] = -1,
    ea: Annotated[str, "Address of the number instead of line/col"] = "",
    opnum: Annotated[int, "Operand number, with ``ea``"] = -1,
) -> dict:
    """Change how a number is displayed in the DECOMPILATION (Hex-Rays keeps its
    own number formats, per (address, operand), independent of the listing).

    Same stops as ``op_format`` minus the two C can't express: binary (no such
    literal -- IDA takes the format and prints decimal anyway) and offset (it
    makes the function stop decompiling). Returns the re-rendered line, and
    marks the function dirty so the next decompile is the new text."""
    import ida_hexrays
    import ida_lines
    import idaapi

    if not ida_hexrays.init_hexrays_plugin():
        return {"addr": str(addr), "error": "no decompiler"}
    try:
        f = idaapi.get_func(parse_address(addr))
    except Exception as e:
        return {"addr": str(addr), "error": str(e)}
    if f is None:
        return {"addr": str(addr), "error": "no function here"}
    try:
        cf = ida_hexrays.decompile(f.start_ea)
    except Exception as e:
        return {"addr": hex(f.start_ea), "error": f"decompile failed: {e}"}
    if cf is None:
        return {"addr": hex(f.start_ea), "error": "decompilation failed"}

    sv = cf.get_pseudocode()
    line = int(line)
    target = None
    if ea:
        try:
            want_ea = parse_address(ea)
        except Exception as e:
            return {"addr": hex(f.start_ea), "error": str(e)}
        for i in range(len(sv)):
            for rec in _idatui_pc_nums(cf, sv[i]):
                if rec["ea"] == want_ea and (int(opnum) < 0
                                             or rec["opnum"] == int(opnum)):
                    target, line = rec, i
                    break
            if target:
                break
    elif 0 <= line < len(sv):
        nums = _idatui_pc_nums(cf, sv[line])
        if nums:
            if int(col) >= 0:
                plain = ida_lines.tag_remove(sv[line].line)
                x = _idatui_uncompact_col(plain, _idatui_compact(plain), int(col))
                target = next((r for r in nums if r["x0"] <= x < r["x1"]), None)
            target = target or nums[0]
    else:
        return {"addr": hex(f.start_ea),
                "error": f"line {line} is outside the {len(sv)}-line decompilation"}
    if target is None:
        return {"addr": hex(f.start_ea), "line": line,
                "text": (ida_lines.tag_remove(sv[line].line).strip()
                         if 0 <= line < len(sv) else ""),
                "error": "no number literal on this line"}

    cur, value = target["fmt"], target["value"]
    choices = [c for c in _IDATUI_PC_FMT_CYCLE
               if c != "char" or _idatui_printable(value)]
    # Same rule as the listing: one ring per literal, every step. A format the
    # ring can't hold (an enum set in the GUI) is reported on the way out
    # instead of being kept for one lap and then lost.
    lossy = cur not in choices and cur != "default"
    out = {"addr": hex(f.start_ea), "ea": hex(target["ea"]),
           "opnum": target["opnum"], "line": line, "prev": cur,
           "format": cur, "shown": target["shown"], "choices": choices,
           "value": hex(value),
           "before": ida_lines.tag_remove(sv[line].line).strip()}

    mode = str(mode or "cycle").lower()
    if mode == "show":
        out["text"] = out["before"]
        out["applied"] = False
        return out
    if mode in ("cycle", "back"):
        step = 1 if mode == "cycle" else -1
        if cur in choices:
            want = choices[(choices.index(cur) + step) % len(choices)]
        else:
            want = choices[0] if step > 0 else choices[-1]
    else:
        want = mode
        if want in ("bin", "offset", "stack", "seg", "float"):
            out["error"] = (f"Hex-Rays has no {want} format for a number "
                            f"-- set it on the listing instead")
            out["text"] = out["before"]
            return out
        if want not in ("hex", "dec", "oct", "char", "default"):
            out["error"] = (f"unknown format {mode!r}; one of hex, dec, oct, "
                            f"char, default")
            out["text"] = out["before"]
            return out

    loc = ida_hexrays.operand_locator_t(target["ea"], target["opnum"])
    it = ida_hexrays.user_numforms_find(cf.numforms, loc)
    if it != ida_hexrays.user_numforms_end(cf.numforms):
        # std::map::insert is a no-op on an existing key, so a format already
        # set here would silently win over the new one.
        ida_hexrays.user_numforms_erase(cf.numforms, it)
    if want != "default":
        import ida_bytes
        nf = ida_hexrays.number_format_t(target["opnum"])
        nf.flags = ida_bytes.get_operand_flag(_idatui_fmt_nibbles()[want],
                                              target["opnum"])
        try:
            nf.org_nbytes = target["nbytes"]
        except Exception:
            pass
        ida_hexrays.user_numforms_insert(cf.numforms, loc, nf)
    cf.save_user_numforms()
    try:
        ida_hexrays.mark_cfunc_dirty(f.start_ea)
    except Exception:
        pass

    out["format"] = want
    out["applied"] = True
    if lossy:
        out["warn"] = (f"this number was {cur}, which names a type a radix "
                       f"can't put back -- reassign it in IDA")
    try:
        cf2 = ida_hexrays.decompile(f.start_ea,
                                    flags=ida_hexrays.DECOMP_NO_CACHE)
        sv2 = cf2.get_pseudocode() if cf2 is not None else None
        out["text"] = (ida_lines.tag_remove(sv2[line].line).strip()
                       if sv2 is not None and line < len(sv2) else out["before"])
    except Exception as e:
        out["text"] = out["before"]
        out["warn"] = f"re-render failed: {e}"
    return out


'''

SNIPPET = f"{BEGIN}\n{BODY.strip()}\n{END}\n"


def api_types_path() -> pathlib.Path | None:
    """Locate ida_pro_mcp/ida_mcp/api_types.py without importing it (importing the
    submodule would pull in IDA, which isn't available outside a worker)."""
    spec = importlib.util.find_spec("ida_pro_mcp")  # top-level pkg is IDA-free
    if spec is None or not spec.submodule_search_locations:
        return None
    p = pathlib.Path(spec.submodule_search_locations[0]) / "ida_mcp" / "api_types.py"
    return p if p.exists() else None


def main() -> int:
    path = api_types_path()
    if path is None:
        print("idatui: ida_pro_mcp not found; skipping tool injection", file=sys.stderr)
        return 0
    text = path.read_text()
    if BEGIN in text and END in text:  # replace the existing block in place
        pre = text[: text.index(BEGIN)].rstrip()
        post = text[text.index(END) + len(END):].lstrip("\n")
        new = pre + "\n\n" + SNIPPET + ("\n" + post if post else "")
    else:
        new = text.rstrip() + "\n\n" + SNIPPET
    if new == text:
        return 0
    try:
        path.write_text(new)
    except OSError as e:
        print(f"idatui: could not patch {path}: {e}", file=sys.stderr)
        return 1
    print(f"idatui: injected/updated idatui-ext tools in {path}", file=sys.stderr)
    return 0


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