#!/usr/bin/env python3
"""Add original-sample UTC time to existing standard scale/power templates.

Dry-run by default. --apply saves only changed templates, verifies readback and
rolls back this invocation's changes on failure. Private complete backups go to
ignored config-backup/. Custom topics/payloads are not rewritten.
"""
import argparse
import json
import os
from pathlib import Path
import re
import time
import urllib.request

TOKEN = re.compile(r'\$\{[^{}]+\}')


class MigrationFailure(RuntimeError):
    def __init__(self, error, rollback):
        super().__init__(str(error))
        self.rollback = rollback


def apply_plans(api, plans, final_verify=lambda: None):
    """Resolve live indices and verify ownership/readback, including lost replies."""
    attempted = []
    def current(plan):
        saved = next(c for c in api('config/backup').get('tcpIo', [])
                     if c['channelId'] == plan['channelId'])
        live = next(c for c in api('tcpio')['channels'] if c['channelId'] == plan['channelId'])
        return saved['converterTemplate'], live['idx']
    try:
        for plan in plans:
            template, idx = current(plan)
            if template != plan['before']:
                raise RuntimeError('Template changed since preflight')
            attempted.append(plan)
            api('converter/template', {'idx': idx, 'template': plan['after']})
            if current(plan)[0] != plan['after']:
                raise RuntimeError('Template readback mismatch')
        final_verify()
    except Exception as error:
        rollback = []
        for plan in reversed(attempted):
            failure = None
            for retry in range(3):
                try:
                    template, idx = current(plan)
                    if template == plan['before']:
                        failure = None; break
                    if template != plan['after']:
                        raise RuntimeError('Template changed by another writer; refusing overwrite')
                    api('converter/template', {'idx': idx, 'template': plan['before']})
                    if current(plan)[0] != plan['before']:
                        raise RuntimeError('Rollback readback mismatch')
                    failure = None; break
                except Exception as e:
                    failure = str(e)
            rollback.append({'channelId': plan['channelId'], 'restored': failure is None, 'error': failure})
        raise MigrationFailure(error, rollback) from error


def upgrade_template(template, topic):
    key = 'weight' if topic.endswith('/scale/weight') else 'v' if topic.endswith('/power') else None
    if not key:
        return template
    tokens = []
    def replace(match):
        tokens.append(match.group())
        return json.dumps('__sample_token_' + str(len(tokens)-1))
    try:
        data = json.loads(TOKEN.sub(replace, template))
    except (ValueError, TypeError) as error:
        raise ValueError('Standard measurement template is not supported JSON') from error
    if not isinstance(data, dict) or key not in data:
        raise ValueError('Standard measurement payload has no primary sample field')
    value = data[key]
    if not isinstance(value, str) or not re.fullmatch(r'__sample_token_\d+', value):
        raise ValueError('Primary sample must be an RS485 token')
    token = tokens[int(value.removeprefix('__sample_token_'))]
    match = re.fullmatch(r'\$\{(rs485\.\d+\.[^{}:]+)(?::\d+)?\}', token)
    if not match or match[1].endswith('.timestamp'):
        raise ValueError('Primary sample must be an RS485 value')
    timestamp = '${' + match[1] + '.timestamp}'
    if 'timestamp' in data:
        value = data['timestamp']
        if isinstance(value, str) and re.fullmatch(r'__sample_token_\d+', value) and tokens[int(value.removeprefix('__sample_token_'))] == timestamp:
            return template
        raise ValueError('Existing timestamp differs; review it before migration')
    trimmed = template.rstrip()
    return trimmed[:-1] + ',"timestamp":' + timestamp + '}' + template[len(trimmed):]


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument('--gw', required=True)
    ap.add_argument('--apply', action='store_true')
    ap.add_argument('--out', required=True)
    ap.add_argument('--backup-dir', default='config-backup/timestamp-migration', help='Private backup directory, relative to the current working directory')
    a = ap.parse_args()
    base = a.gw if a.gw.startswith('http') else 'http://' + a.gw
    token = os.environ.get('OPTA_TOKEN', 'smmsadmin')
    def api(path, body=None):
        req = urllib.request.Request(base + '/api/' + path, data=None if body is None else json.dumps(body,ensure_ascii=False,separators=(',',':')).encode(), headers={'Authorization':'Bearer '+token,'Content-Type':'application/json','Connection':'close'})
        with urllib.request.urlopen(req, timeout=20) as response:
            return json.load(response)
    diag = api('diag')
    version = tuple(int(x) for x in diag['fwVersion'].split('.'))
    if version < (6,0,120):
        raise RuntimeError('Firmware must support original sample timestamps (>=6.0.120)')
    before = api('config/backup')
    channels = api('tcpio')['channels']
    plans = []
    for saved in before.get('tcpIo', []):
        if saved.get('mode') != 2:
            continue
        old = saved.get('converterTemplate', '')
        new = upgrade_template(old, saved.get('mqttTopic', ''))
        if old != new:
            live = next(c for c in channels if c['channelId'] == saved['channelId'])
            plans.append({'channelId':saved['channelId'],'idx':live['idx'],'before':old,'after':new})
    output = Path(a.out);output.parent.mkdir(parents=True,exist_ok=True)
    private = Path(a.backup_dir).expanduser().resolve()
    try:
        if a.apply:
            private.mkdir(parents=True,exist_ok=True)
            (private/(str(time.time_ns())+'.json')).write_text(json.dumps(before,ensure_ascii=False))
            def verify_boot():
                assert api('diag')['bootCount']==diag['bootCount'], 'Unexpected reboot'
            apply_plans(api, plans, verify_boot)
        output.write_text(json.dumps({'applied':a.apply,'fw':diag['fwVersion'],'buildId':diag['fault']['currentBuildId'],'changes':plans},ensure_ascii=False,indent=2)+'\n')
        print(f'{"Applied" if a.apply else "Dry run"}: {len(plans)} standard measurement templates')
    except Exception as error:
        output.write_text(json.dumps({'applied':False,'pass':False,'error':str(error),
            'rollback':getattr(error,'rollback',[]),'changes':plans},ensure_ascii=False,indent=2)+'\n')
        raise


if __name__ == '__main__':
    main()