#!/usr/bin/env python3
import argparse
import json
import shutil
import subprocess
import time
from pathlib import Path
from typing import List, Optional

OPENCLAW_HOME = Path.home() / '.openclaw'
AGENTS_DIR = OPENCLAW_HOME / 'agents'
CONFIG_PATH = OPENCLAW_HOME / 'openclaw.json'
WEIXIN_DIR = OPENCLAW_HOME / 'openclaw-weixin'
WEIXIN_ACCOUNTS_DIR = WEIXIN_DIR / 'accounts'
WEIXIN_CREDENTIALS_DIR = OPENCLAW_HOME / 'credentials'


def load_json(path: Path):
    with path.open() as f:
        return json.load(f)


def save_json(path: Path, data):
    tmp = path.with_suffix(path.suffix + '.tmp')
    with tmp.open('w') as f:
        json.dump(data, f, ensure_ascii=False, indent=2)
        f.write('\n')
    tmp.replace(path)


def normalize(s: Optional[str]) -> str:
    return (s or '').lower()


def is_legacy_key(key: str) -> bool:
    parts = key.split(':')
    return len(parts) >= 5 and parts[:3] == ['agent', parts[1], 'openclaw-weixin'] and parts[3] == 'direct'


def match_entry(key: str, entry: dict, channel: str, account_id: str, peer: Optional[str]) -> bool:
    origin = entry.get('origin') or {}
    delivery = entry.get('deliveryContext') or {}
    last_account = entry.get('lastAccountId')
    key_n = normalize(key)
    peer_n = normalize(peer)
    account_n = normalize(account_id)
    channel_n = normalize(channel)
    account_hit = (
        normalize(origin.get('accountId')) == account_n
        or normalize(delivery.get('accountId')) == account_n
        or normalize(last_account) == account_n
        or account_n in key_n
    )
    channel_hit = (
        normalize(origin.get('provider')) == channel_n
        or normalize(delivery.get('channel')) == channel_n
        or channel_n in key_n
    )
    peer_hit = True if not peer else peer_n in key_n or normalize(origin.get('from')) == peer_n or normalize(delivery.get('to')) == peer_n
    return account_hit and channel_hit and peer_hit


def find_matches(channel: str, account_id: str, peer: Optional[str]) -> List[dict]:
    matches = []
    for agent_dir in sorted(AGENTS_DIR.iterdir()):
        sessions_path = agent_dir / 'sessions' / 'sessions.json'
        if not sessions_path.exists():
            continue
        data = load_json(sessions_path)
        for key, entry in data.items():
            if match_entry(key, entry, channel, account_id, peer):
                matches.append({
                    'agent': agent_dir.name,
                    'sessionsPath': sessions_path,
                    'data': data,
                    'sessionKey': key,
                    'entry': entry,
                    'legacyKey': is_legacy_key(key),
                })
    return matches


def backup_and_remove(match: dict, dry_run: bool):
    sessions_path = match['sessionsPath']
    data = match['data']
    session_key = match['sessionKey']
    entry = match['entry']
    base = sessions_path.parent
    backup_dir = base / 'manual-backups'
    ts = time.strftime('%Y%m%d-%H%M%S')
    session_file = Path(entry.get('sessionFile', '')) if entry.get('sessionFile') else None
    result = {
        'agent': match['agent'],
        'sessionKey': session_key,
        'legacyKey': match['legacyKey'],
        'entryBackup': str(backup_dir / f'{ts}-session-entry-{match["agent"]}.json'),
        'transcriptBackup': None,
        'removed': False,
    }
    if session_file and session_file.exists():
        result['transcriptBackup'] = str(backup_dir / f'{ts}-{session_file.name}')
    if dry_run:
        return result
    backup_dir.mkdir(exist_ok=True)
    Path(result['entryBackup']).write_text(json.dumps({session_key: entry}, ensure_ascii=False, indent=2))
    if session_file and session_file.exists():
        shutil.copy2(session_file, Path(result['transcriptBackup']))
    data.pop(session_key, None)
    save_json(sessions_path, data)
    result['removed'] = True
    return result


def clear_bindings(channel: str, account_id: str, dry_run: bool):
    cfg = load_json(CONFIG_PATH)
    bindings = cfg.get('bindings', [])
    account_n = normalize(account_id)
    channel_n = normalize(channel)
    kept = []
    removed = []
    for b in bindings:
        match = b.get('match') or {}
        same = normalize(match.get('channel')) == channel_n and normalize(match.get('accountId')) == account_n
        if same:
            removed.append(b)
        else:
            kept.append(b)
    if removed and not dry_run:
        cfg['bindings'] = kept
        backup = CONFIG_PATH.with_suffix('.json.bak')
        shutil.copy2(CONFIG_PATH, backup)
        save_json(CONFIG_PATH, cfg)
    return removed


def clear_weixin_account_state(account_id: str, dry_run: bool):
    targets = [
        WEIXIN_ACCOUNTS_DIR / f'{account_id}.json',
        WEIXIN_ACCOUNTS_DIR / f'{account_id}.sync.json',
        WEIXIN_ACCOUNTS_DIR / f'{account_id}.context-tokens.json',
        WEIXIN_CREDENTIALS_DIR / f'openclaw-weixin-{account_id}-allowFrom.json',
    ]
    backup_dir = WEIXIN_DIR / 'manual-backups'
    ts = time.strftime('%Y%m%d-%H%M%S')
    results = []
    for path in targets:
        item = {'path': str(path), 'exists': path.exists(), 'backup': None, 'removed': False}
        if path.exists():
            item['backup'] = str(backup_dir / f'{ts}-{path.name}')
        if dry_run:
            results.append(item)
            continue
        if path.exists():
            backup_dir.mkdir(parents=True, exist_ok=True)
            shutil.copy2(path, item['backup'])
            path.unlink()
            item['removed'] = True
        results.append(item)
    return results


def restart_gateway(dry_run: bool):
    if dry_run:
        return {'requested': True, 'executed': False}
    cp = subprocess.run(['openclaw', 'gateway', 'start'], capture_output=True, text=True)
    return {'requested': True, 'executed': cp.returncode == 0, 'stdout': cp.stdout.strip(), 'stderr': cp.stderr.strip()}


def main():
    ap = argparse.ArgumentParser(description='Reset a DM/channel account by clearing live sessions, bindings, and optionally weixin account state. Use --peer for soft rebind (keeps account state).')
    ap.add_argument('--channel', default='openclaw-weixin')
    ap.add_argument('--account-id', required=True)
    ap.add_argument('--peer')
    ap.add_argument('--restart-gateway', action='store_true')
    ap.add_argument('--yes', action='store_true')
    ap.add_argument('--dry-run', action='store_true')
    args = ap.parse_args()

    dry_run = args.dry_run or not args.yes
    soft = bool(args.peer)  # --peer means soft rebind: keep account state
    matches = find_matches(args.channel, args.account_id, args.peer)
    report = {
        'mode': 'reset',
        'channel': args.channel,
        'accountId': args.account_id,
        'peer': args.peer,
        'soft': soft,
        'dryRun': dry_run,
        'matches': [
            {
                'agent': m['agent'],
                'sessionKey': m['sessionKey'],
                'legacyKey': m['legacyKey'],
                'sessionId': m['entry'].get('sessionId'),
                'sessionFile': m['entry'].get('sessionFile'),
            }
            for m in matches
        ]
    }
    print(json.dumps(report, ensure_ascii=False, indent=2))

    actions = [backup_and_remove(m, dry_run) for m in matches]
    # soft mode (--peer): only clear sessions, keep binding and account state
    removed_bindings = clear_bindings(args.channel, args.account_id, dry_run) if not soft else []
    weixin_state = None if soft else clear_weixin_account_state(args.account_id, dry_run)
    gateway = restart_gateway(dry_run) if args.restart_gateway else None
    result = {
        'removedSessions': actions,
        'bindingsRemoved': removed_bindings,
        'weixinAccountState': weixin_state,
        'softRebind': soft,
        'gatewayRestart': gateway,
        'nextStep': 'Send a fresh inbound message to verify the new routing, or re-scan/login the weixin account if doing a full reset.'
    }
    print(json.dumps(result, ensure_ascii=False, indent=2))


if __name__ == '__main__':
    main()
