#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
st-card-tool.py —— 酒馆（SillyTavern）角色卡「验卡 / 救卡 / 批量导入」工具
适用：本机 DSH 机上的酒馆 http://127.0.0.1:8000

为什么需要它：
  1) 酒馆只认 PNG（内嵌 tEXt 块 chara/ccv3）与 json/yaml/charx/byaf，**完全不支持 webp**；
  2) 社区（如类脑）页面上的 webp 往往是**纯封面图**，转成 PNG 也依然没有角色数据；
  3) 本工具能判别"真卡 vs 封面"，并把带数据的 webp/图片卡救成酒馆能用的 PNG。

用法：
  python3 st-card-tool.py inspect <文件或目录...>     # 验卡：是不是真卡、角色名、规格版本
  python3 st-card-tool.py import  <文件或目录...>     # 批量导入酒馆（自动救 webp、自动选 file_type）
  python3 st-card-tool.py list                        # 列出酒馆里已有的角色
  python3 st-card-tool.py list --detail               # 附带头像/世界书/备用开场等信息

常用参数：
  --url http://127.0.0.1:8000   酒馆地址（默认本机）
  --out DIR                     救卡产物输出目录（默认 ./st-cards-converted）
  --dry-run                     导入前只报告，不真正写库
"""
import argparse
import base64
import json
import os
import re
import shutil
import struct
import subprocess
import sys
import tempfile
import urllib.request
import uuid
import zlib
from pathlib import Path

ST_URL = 'http://127.0.0.1:8000'
CHARS_DIR = Path('/home/zyw/SillyTavern/data/default-user/characters')

CARD_EXTS = {'.png', '.webp', '.json', '.yaml', '.yml', '.charx', '.byaf', '.jpg', '.jpeg'}
# 酒馆 /api/characters/import 支持的 file_type
ST_FILE_TYPES = {'png': 'png', 'json': 'json', 'yaml': 'yaml', 'yml': 'yml',
                 'charx': 'charx', 'byaf': 'byaf'}

G, Y, R, C, B, N = '\033[32m', '\033[33m', '\033[31m', '\033[36m', '\033[1m', '\033[0m'


def ok(s):    return f'{G}✅ {s}{N}'
def warn(s):  return f'{Y}⚠️  {s}{N}'
def bad(s):   return f'{R}❌ {s}{N}'
def info(s):  return f'{C}ℹ️  {s}{N}'


# ─────────────────────────── 图片/容器格式探测 ───────────────────────────

def sniff(path: Path) -> str:
    head = path.open('rb').read(16)
    if head.startswith(b'\x89PNG\r\n\x1a\n'):
        return 'png'
    if head.startswith(b'\xff\xd8\xff'):
        return 'jpeg'
    if head[:4] == b'RIFF' and head[8:12] == b'WEBP':
        return 'webp'
    if head[:4] == b'PK\x03\x04':
        return 'zip'          # charx = zip；byaf 也可能
    if head[:2] == b'\x1f\x8b':
        return 'gzip'
    try:
        txt = path.open('rb').read(4096).decode('utf-8', 'ignore').lstrip()
        if txt.startswith('{'):
            return 'json'
    except Exception:
        pass
    return 'unknown'


# ─────────────────────── PNG：chara / ccv3 tEXt 块 ───────────────────────

def png_chunks(buf: bytes):
    """迭代 PNG chunk：返回 (type, data, offset)"""
    if buf[:8] != b'\x89PNG\r\n\x1a\n':
        return
    i = 8
    while i + 8 <= len(buf):
        (ln,) = struct.unpack('>I', buf[i:i + 4])
        ctype = buf[i + 4:i + 8]
        data = buf[i + 8:i + 8 + ln]
        yield ctype, data, i
        i += 12 + ln
        if ctype == b'IEND':
            break


def png_read_text_chunks(buf: bytes) -> dict:
    """返回 {keyword: 原始文本}（tEXt 与 iTXt 都读）"""
    out = {}
    for ctype, data, _ in png_chunks(buf) or []:
        if ctype == b'tEXt':
            kw, _, val = data.partition(b'\x00')
            out[kw.decode('latin1')] = val.decode('latin1')
        elif ctype == b'iTXt':
            kw, _, rest = data.partition(b'\x00')
            # compression flag(1) method(1) lang\0 translated\0 text
            try:
                rest = rest[2:]
                _lang, _, rest = rest.partition(b'\x00')
                _trans, _, text = rest.partition(b'\x00')
                out[kw.decode('latin1')] = text.decode('utf-8', 'ignore')
            except Exception:
                pass
    return out


def png_write_with_chara(base_png: bytes, card: dict) -> bytes:
    """把 card JSON 作为 tEXt 块 (chara + ccv3) 写进 PNG（替换已有的），返回新 PNG"""
    chunks = [c for c in (png_chunks(base_png) or [])]
    if not chunks:
        raise ValueError('不是合法的 PNG')

    keep = []
    for ctype, data, _ in chunks:
        if ctype == b'tEXt':
            kw = data.partition(b'\x00')[0].decode('latin1').lower()
            if kw in ('chara', 'ccv3'):
                continue
        keep.append((ctype, data))

    payload = json.dumps(card, ensure_ascii=False)

    def text_chunk(kw: str, text: str) -> bytes:
        body = kw.encode('latin1') + b'\x00' + text.encode('utf-8')
        return (struct.pack('>I', len(body)) + b'tEXt' + body
                + struct.pack('>I', zlib.crc32(b'tEXt' + body) & 0xffffffff))

    # 插到 IEND 之前
    out = bytearray(b'\x89PNG\r\n\x1a\n')
    for ctype, data in keep:
        if ctype == b'IEND':
            break
        out += struct.pack('>I', len(data)) + ctype + data
        out += struct.pack('>I', zlib.crc32(ctype + data) & 0xffffffff)
    # v2 chara
    out += text_chunk('chara', base64.b64encode(payload.encode('utf-8')).decode('ascii'))
    # v3 ccv3（若卡本体是 v2，顺手升级一份）
    v3 = dict(card)
    v3['spec'] = 'chara_card_v3'
    v3['spec_version'] = '3.0'
    out += text_chunk('ccv3', base64.b64encode(json.dumps(v3, ensure_ascii=False).encode('utf-8')).decode('ascii'))
    # IEND
    out += struct.pack('>I', 0) + b'IEND' + struct.pack('>I', zlib.crc32(b'IEND') & 0xffffffff)
    return bytes(out)


# ───────────────────────── WebP：RIFF 块扫描 ─────────────────────────

def webp_chunks(buf: bytes):
    if buf[:4] != b'RIFF' or buf[8:12] != b'WEBP':
        return
    i = 12
    while i + 8 <= len(buf):
        fourcc = buf[i:i + 4]
        (ln,) = struct.unpack('<I', buf[i + 4:i + 8])
        data = buf[i + 8:i + 8 + ln]
        yield fourcc.decode('latin1'), data
        i += 8 + ln + (ln & 1)


# ─────────────────── 角色数据抽取（JSON / base64 / V1~V3） ───────────────────

def _looks_like_card(d) -> bool:
    if not isinstance(d, dict):
        return False
    if isinstance(d.get('data'), dict) and ('name' in d['data'] or 'description' in d['data']):
        return True
    return 'name' in d and ('description' in d or 'first_mes' in d or 'personality' in d)


def _brace_json_candidates(text: str):
    """从任意文本里用花括号配对抠出 JSON 候选"""
    for m in re.finditer(r'\{', text):
        depth, instr, esc = 0, False, False
        for i in range(m.start(), len(text)):
            ch = text[i]
            if instr:
                if esc:
                    esc = False
                elif ch == '\\':
                    esc = True
                elif ch == '"':
                    instr = False
                continue
            if ch == '"':
                instr = True
            elif ch == '{':
                depth += 1
            elif ch == '}':
                depth -= 1
                if depth == 0:
                    yield text[m.start():i + 1]
                    break


def extract_card_from_bytes(buf: bytes):
    """尽最大努力从任意字节流里找出角色卡 JSON。返回 (card, 来源说明) 或 (None, None)"""
    # 1) 直接就是 JSON 文本
    try:
        d = json.loads(buf.decode('utf-8'))
        if _looks_like_card(d):
            return d, 'plain-json'
    except Exception:
        pass

    # 2) PNG tEXt / iTXt 的 chara / ccv3
    try:
        tc = png_read_text_chunks(buf)
        for kw in ('ccv3', 'chara'):
            if kw in tc:
                try:
                    d = json.loads(base64.b64decode(tc[kw]).decode('utf-8'))
                    if _looks_like_card(d):
                        return d, f'png-{kw}'
                except Exception:
                    pass
    except Exception:
        pass

    # 3) 文本里直接出现的 JSON
    text = buf.decode('utf-8', 'ignore')
    for cand in _brace_json_candidates(text):
        try:
            d = json.loads(cand)
            if _looks_like_card(d):
                return d, 'embedded-json'
        except Exception:
            continue

    # 4) base64 块（chara 常见于图片元数据）
    for m in re.finditer(rb'[A-Za-z0-9+/=]{200,}', buf):
        blob = m.group()
        for trim in (0, 1, 2):
            try:
                raw = base64.b64decode(blob[:len(blob) - trim] if trim else blob)
                d = json.loads(raw.decode('utf-8'))
                if _looks_like_card(d):
                    return d, 'base64-blob'
            except Exception:
                continue
    return None, None


def card_name(card: dict) -> str:
    if isinstance(card.get('data'), dict) and card['data'].get('name'):
        return str(card['data']['name'])
    return str(card.get('name') or '(未命名字符卡)')


def card_meta(card: dict) -> dict:
    d = card.get('data') if isinstance(card.get('data'), dict) else card
    return {
        'name': card_name(card),
        'spec': card.get('spec') or ('chara_card_v1(无 spec)' if 'name' in card else '?'),
        'spec_version': card.get('spec_version') or '-',
        'creator': d.get('creator') or d.get('created_by') or '-',
        'version': d.get('character_version') or '-',
        'tags': d.get('tags') or [],
        'desc_len': len(str(d.get('description') or '')),
        'first_mes': str(d.get('first_mes') or ''),
        'alt_greetings': len(d.get('alternate_greetings') or []),
        'has_world': bool((d.get('character_book') or {}).get('entries')),
        'world_entries': len(((d.get('character_book') or {}).get('entries')) or []),
    }


# ─────────────────────────── inspect ───────────────────────────

def iter_files(paths):
    for p in paths:
        p = Path(p)
        if p.is_dir():
            for f in sorted(p.rglob('*')):
                if f.is_file() and f.suffix.lower() in CARD_EXTS:
                    yield f
        elif p.is_file():
            yield p
        else:
            print(warn(f'找不到: {p}'))


def cmd_inspect(args):
    files = list(iter_files(args.paths))
    if not files:
        print(bad('没有可检查的文件')); return 1
    real = fake = 0
    for f in files:
        kind = sniff(f)
        size = f.stat().st_size
        card, src = extract_card_from_bytes(f.read_bytes())

        if card:
            m = card_meta(card)
            real += 1
            print(ok(f'{f.name}  →  {B}真卡{N}'))
            print(f'    格式 {kind} / {size/1024:.1f} KB   数据来源 {src}')
            print(f'    角色名 {B}{m["name"]}{N}   规格 {m["spec"]} {m["spec_version"]}')
            print(f'    描述 {m["desc_len"]} 字   备用开场 {m["alt_greetings"]} 条   '
                  f'世界书 {"有(" + str(m["world_entries"]) + "条)" if m["has_world"] else "无"}')
        else:
            fake += 1
            extra = ''
            if kind == 'webp':
                ch = [c for c, _ in (webp_chunks(f.read_bytes()) or [])]
                extra = f'   RIFF 块: {", ".join(ch) or "无"}'
            print(bad(f'{f.name}  →  {B}不含角色数据{N}（只是封面图）'))
            print(f'    格式 {kind} / {size/1024:.1f} KB{extra}')
            print(f'    {Y}请回到卡片页面点「下载角色卡 / 导出 PNG / 下载 JSON」，别右键"图片另存为"{N}')

    print()
    print(f'汇总：真卡 {G}{real}{N} 个，无数据 {R}{fake}{N} 个，共 {len(files)} 个')
    return 0


# ─────────────────────────── 救卡：转成酒馆可用 ───────────────────────────

def convert_one(f: Path, outdir: Path):
    """返回 (产出文件 Path, file_type) 或 (None, None, 原因)"""
    kind = sniff(f)
    raw = f.read_bytes()

    # 已经是酒馆能直接吃的格式
    if kind == 'json' and f.suffix.lower() == '.json':
        card, _ = extract_card_from_bytes(raw)
        return f, 'json' if card else None
    if kind == 'zip':
        return f, f.suffix.lower().lstrip('.')     # charx / byaf
    if kind == 'png':
        card, _ = extract_card_from_bytes(raw)
        if card:
            return f, 'png'
        # 是 PNG 但没有数据 —— 也救不了
        return None, None

    # webp / jpeg / 其它图片：找数据 → 转 PNG → 注入 chara
    card, src = extract_card_from_bytes(raw)
    if not card:
        return None, None

    outdir.mkdir(parents=True, exist_ok=True)
    safe = re.sub(r'[\\/:*?"<>|\x00-\x1f]', '_', card_name(card))[:80] or 'card'
    out_png = outdir / f'{safe}.png'

    # 用 ffmpeg 把图片转成 PNG（webp/jpeg/… 都能解）
    tmpdir = Path(tempfile.mkdtemp(prefix='stcard-'))
    try:
        src_img = tmpdir / f'img{f.suffix or ".bin"}'
        src_img.write_bytes(raw)
        plain = tmpdir / 'plain.png'
        r = subprocess.run(['ffmpeg', '-y', '-loglevel', 'error', '-i', str(src_img), str(plain)],
                           capture_output=True, text=True)
        if r.returncode != 0 or not plain.exists():
            return None, None
        out_png.write_bytes(png_write_with_chara(plain.read_bytes(), card))
    finally:
        shutil.rmtree(tmpdir, ignore_errors=True)
    return out_png, 'png'


# ─────────────────────────── 调用酒馆 API ───────────────────────────

class ST:
    def __init__(self, url):
        self.url = url.rstrip('/')
        self.cookie = None
        self.csrf = None

    def _opener(self):
        return urllib.request.build_opener(urllib.request.HTTPCookieProcessor())

    def connect(self):
        op = self._opener()
        with op.open(f'{self.url}/csrf-token', timeout=15) as r:
            self.csrf = json.loads(r.read())['token']
        self.opener = op
        return self

    def import_file(self, path: Path, file_type: str):
        boundary = '----stcard' + uuid.uuid4().hex
        body = b''
        body += f'--{boundary}\r\n'.encode()
        body += (f'Content-Disposition: form-data; name="avatar"; filename="{path.name}"\r\n'
                 f'Content-Type: application/octet-stream\r\n\r\n').encode()
        body += path.read_bytes() + b'\r\n'
        body += f'--{boundary}\r\n'.encode()
        body += b'Content-Disposition: form-data; name="file_type"\r\n\r\n'
        body += file_type.encode() + b'\r\n'
        body += f'--{boundary}--\r\n'.encode()

        req = urllib.request.Request(
            f'{self.url}/api/characters/import', data=body, method='POST',
            headers={'Content-Type': f'multipart/form-data; boundary={boundary}',
                     'X-CSRF-Token': self.csrf, 'Content-Length': str(len(body))})
        try:
            with self.opener.open(req, timeout=120) as r:
                return r.status, json.loads(r.read())
        except urllib.error.HTTPError as e:
            return e.code, e.read().decode('utf-8', 'ignore')[:300]


def cmd_import(args):
    files = list(iter_files(args.paths))
    if not files:
        print(bad('没有可导入的文件')); return 1

    outdir = Path(args.out)
    st = ST(args.url).connect()
    done = skipped = failed = 0

    for f in files:
        kind = sniff(f)
        # 先判断是不是真卡（json/zip 直接放行）
        card, _ = extract_card_from_bytes(f.read_bytes()) if kind not in ('zip',) else (None, None)
        if kind != 'zip' and not card:
            print(bad(f'{f.name}：不含角色数据（只是封面），跳过'))
            skipped += 1
            continue

        if kind in ('webp', 'jpeg', 'unknown'):
            out, ftype = convert_one(f, outdir)
            if not out:
                print(bad(f'{f.name}：救卡失败（ffmpeg 解不开或没有图片数据），跳过'))
                failed += 1
                continue
            print(info(f'{f.name}：已救成酒馆可用 PNG → {out}'))
            target = out
        else:
            ext = f.suffix.lower().lstrip('.')
            ftype = ST_FILE_TYPES.get(ext)
            target = f
            if not ftype:
                print(bad(f'{f.name}：酒馆不支持 .{ext}，跳过')); skipped += 1; continue

        if args.dry_run:
            print(info(f'[dry-run] 会导入 {target.name} (file_type={ftype})'))
            continue

        code, resp = st.import_file(target, ftype)
        if code == 200 and isinstance(resp, dict) and resp.get('file_name'):
            print(ok(f'已导入：{B}{resp["file_name"]}{N}   (源 {f.name})'))
            done += 1
        else:
            print(bad(f'{f.name} 导入失败 HTTP {code}: {resp}'))
            failed += 1

    print()
    print(f'汇总：导入成功 {G}{done}{N}，跳过 {Y}{skipped}{N}，失败 {R}{failed}{N}')
    if done and not args.dry_run:
        print(info('回到酒馆页面按 Ctrl+F5 刷新，左侧「角色管理」里就能看到新卡了'))
    return 0


# ─────────────────────────── list ───────────────────────────

def cmd_list(args):
    if not CHARS_DIR.exists():
        print(bad(f'找不到角色目录 {CHARS_DIR}')); return 1
    rows = []
    for f in sorted(CHARS_DIR.glob('*.png')):
        card, _ = extract_card_from_bytes(f.read_bytes())
        if card:
            m = card_meta(card)
            rows.append((m['name'], f.name, f.stat().st_size,
                         m['alt_greetings'], m['world_entries'], m['spec']))
        else:
            rows.append((f.stem, f.name, f.stat().st_size, '-', '-', '无数据'))
    if not rows:
        print(warn('还没有任何角色卡')); return 0
    print(f'酒馆里共 {B}{len(rows)}{N} 个角色（{CHARS_DIR}）\n')
    for name, fn, size, alt, wi, spec in rows:
        if args.detail:
            print(f'  {B}{name}{N}')
            print(f'      文件 {fn}  {size/1024:.0f} KB   规格 {spec}   '
                  f'备用开场 {alt}   世界书条目 {wi}')
        else:
            print(f'  {name}')
    return 0


# ─────────────────────────── main ───────────────────────────

def main():
    ap = argparse.ArgumentParser(description='酒馆角色卡 验卡/救卡/批量导入 工具',
                                 formatter_class=argparse.RawDescriptionHelpFormatter,
                                 epilog=__doc__)
    sub = ap.add_subparsers(dest='cmd')

    for name, fn in (('inspect', cmd_inspect), ('import', cmd_import)):
        p = sub.add_parser(name)
        p.add_argument('paths', nargs='+', help='文件或目录（目录会递归找卡）')
        p.add_argument('--url', default=ST_URL)
        if name == 'import':
            p.add_argument('--out', default='./st-cards-converted', help='救卡产物输出目录')
            p.add_argument('--dry-run', action='store_true')
            p.set_defaults(func=fn)
        else:
            p.set_defaults(func=fn)

    p = sub.add_parser('list')
    p.add_argument('--detail', action='store_true')
    p.set_defaults(func=cmd_list)

    args = ap.parse_args()
    if not args.cmd:
        ap.print_help(); return 0
    return args.func(args)


if __name__ == '__main__':
    sys.exit(main())