#!/usr/bin/env python3
# SPDX-License-Identifier: GPL-3.0-or-later
# Copyright (C) 2025 - 2026, Patrick Schurig / ETM PowerSync
"""Rejoue un fichier de traces de décision vers un optimiseur (OPTIMIZER_PROTOCOL.md annexe D).

    etm-powersync-trace-replay TRACE.ndjson [--endpoint unix:///run/powersync/optimizer.sock]
                               [--token-file FICHIER] [--out SORTIE.ndjson] [--timeout S]
                               [--tolerance-w W] [--fail-on-divergence]

Chaque ligne de TRACE qui porte une `request` est renvoyée, telle quelle, en `GetPlan` à
l'optimiseur (après un `Handshake`). Les réponses sont écrites dans un NOUVEAU fichier, à côté
de la trace (jamais d'écrasement), puis un résumé : requêtes, taux d'abstention, `computeMs`
médian et max, cycles dont les actions diffèrent du plan enregistré.

« Actions du plan » : celles du créneau qui couvre l'instant de la requête (`timestamp`) —
charge, kind, puissance d'une consigne (`setpoint`, à --tolerance-w près), état d'une action
`state` (SG-Ready). La puissance d'une action `state` n'est pas comparée : le protocole ne la
définit pas (§6), et le plan du rule-based la porte (puissance estimée) quand un optimiseur ne la
donne pas. Une abstention (`abstain`, aucun
créneau, aucun créneau courant) n'a pas d'actions. Une ligne dont la réponse enregistrée est une
erreur (`no-response`, `timeout`…) n'est pas comparable : comptée à part.

Le token se lit dans un FICHIER (--token-file), jamais en argument : une ligne de commande se
voit dans `ps` et dans l'historique du shell. Il n'est jamais affiché.

Les traces contiennent des mesures du site : elles restent sur le site. Ce script n'envoie rien
ailleurs qu'à l'endpoint donné.

Code de sortie : 0 ; 1 s'il y a divergence ET --fail-on-divergence ; 2 sur erreur (fichier,
connexion, handshake refusé).
"""

import argparse
import json
import os
import socket
import statistics
import sys
import time

ENDPOINT_DEFAUT = 'unix:///run/powersync/optimizer.sock'
MAJEURE = '1'
VERSION = '1.1'


class Erreur(Exception):
    """Erreur fatale : message pour l'utilisateur, code de sortie 2."""


def ouvrir_endpoint(endpoint, timeout):
    if endpoint.startswith('unix://'):
        chemin = endpoint[len('unix://'):]
        if not chemin.startswith('/'):
            raise Erreur(f'endpoint unix sans chemin absolu : {endpoint}')
        s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        cible = chemin
    elif endpoint.startswith('tcp://'):
        reste = endpoint[len('tcp://'):].rstrip('/')
        if reste.startswith('['):                       # IPv6 entre crochets
            hote, _, port = reste[1:].partition(']:')
        else:
            hote, _, port = reste.rpartition(':')
        if not hote or not port.isdigit():
            raise Erreur(f'endpoint tcp illisible (tcp://hôte:port attendu) : {endpoint}')
        famille = socket.AF_INET6 if ':' in hote else socket.AF_INET
        s = socket.socket(famille, socket.SOCK_STREAM)
        cible = (hote, int(port))
    else:
        raise Erreur(f'endpoint inconnu (unix:///chemin ou tcp://hôte:port) : {endpoint}')
    s.settimeout(timeout)
    try:
        s.connect(cible)
    except OSError as e:
        s.close()
        raise Erreur(f'connexion impossible à {endpoint} : {e.strerror or e}')
    return s


class Client:
    """JSON-RPC 2.0 sur NDJSON, une requête à la fois."""

    def __init__(self, sock):
        self.sock = sock
        self.tampon = b''
        self.id = 0

    def appeler(self, methode, params):
        self.id += 1
        envoi = json.dumps({'jsonrpc': '2.0', 'id': self.id, 'method': methode, 'params': params},
                           separators=(',', ':'), ensure_ascii=False).encode() + b'\n'
        self.sock.sendall(envoi)
        while True:
            message = self.lire()
            if isinstance(message, dict) and message.get('id') == self.id:
                return message
            # notification (UpdateLearned…) ou réponse étrangère : ignorée

    def lire(self):
        while b'\n' not in self.tampon:
            morceau = self.sock.recv(65536)
            if not morceau:
                raise ConnectionError('connexion fermée par l\'optimiseur')
            self.tampon += morceau
        ligne, _, self.tampon = self.tampon.partition(b'\n')
        try:
            return json.loads(ligne)
        except ValueError:
            return None


def lire_token(chemin):
    if not chemin:
        return None
    try:
        with open(chemin, 'rb') as f:
            token = f.read(8192).strip()
    except OSError as e:
        raise Erreur(f'fichier de token illisible : {chemin} : {e.strerror}')
    return token.decode('utf-8', 'replace') or None


def handshake(client, token):
    params = {'protocolVersion': VERSION, 'engine': 'etm-powersync-trace-replay'}
    if token:
        params['token'] = token
    r = client.appeler('Handshake', params)
    if 'error' in r:
        # Seul le CODE : un serveur peut recopier le token refusé dans son message.
        raise Erreur(f'handshake refusé (erreur JSON-RPC {r["error"].get("code")})')
    res = r.get('result')
    if not isinstance(res, dict) or not isinstance(res.get('protocolVersion'), str):
        raise Erreur('handshake : réponse sans protocolVersion')
    if res['protocolVersion'].split('.')[0] != MAJEURE:
        raise Erreur(f'handshake : optimiseur en {res["protocolVersion"][:32]}, majeure {MAJEURE} attendue')
    return res


def instant(texte):
    from datetime import datetime
    if not isinstance(texte, str):
        return None
    try:
        return datetime.fromisoformat(texte.replace('Z', '+00:00'))
    except ValueError:
        return None


def actions_courantes(reponse, maintenant):
    """None si la réponse n'est pas comparable (erreur) ; [] pour une abstention ; sinon la
    liste triée des actions du créneau couvrant `maintenant`."""
    if not isinstance(reponse, dict) or 'result' not in reponse:
        return None
    res = reponse['result']
    if not isinstance(res, dict) or res.get('abstain') is True:
        return []
    for creneau in res.get('slots') or []:
        debut, fin = instant(creneau.get('from')), instant(creneau.get('to'))
        if maintenant is None or debut is None or fin is None or not (debut <= maintenant < fin):
            continue
        actions = []
        for a in creneau.get('actions') or []:
            kind = str(a.get('kind', 'setpoint'))
            actions.append((str(a.get('loadId')), kind,
                            None if kind == 'state' else a.get('powerW'), a.get('state')))
        return sorted(actions, key=lambda x: (x[0], x[1]))
    return []


def differences(attendu, obtenu, tolerance):
    """Liste lisible des écarts entre deux listes d'actions (vide = identiques)."""
    if attendu == [] and obtenu == []:
        return []
    if (attendu == []) != (obtenu == []):
        return [('abstention' if attendu == [] else 'plan') + ' → '
                + ('abstention' if obtenu == [] else 'plan')]
    da = {(a[0], a[1]): a for a in attendu}
    do = {(o[0], o[1]): o for o in obtenu}
    ecarts = []
    for cle in sorted(set(da) | set(do)):
        a, o = da.get(cle), do.get(cle)
        if a is None or o is None:
            ecarts.append(f'{cle[0]} : {"absente" if a is None else "présente"} → '
                          f'{"absente" if o is None else "présente"}')
            continue
        pa, po = a[2], o[2]
        if (pa is None) != (po is None) or (pa is not None and abs(float(pa) - float(po)) > tolerance):
            ecarts.append(f'{cle[0]} : {pa} W → {po} W')
        if a[3] != o[3]:
            ecarts.append(f'{cle[0]} : état {a[3]} → {o[3]}')
    return ecarts


def sortie_libre(trace):
    base = trace[:-len('.ndjson')] if trace.endswith('.ndjson') else trace
    horodatage = time.strftime('%Y%m%dT%H%M%S')
    # Deux rejeux dans la même seconde : suffixe, jamais d'écrasement (ni des réponses, ni du résumé).
    for n in range(1, 1000):
        nom = f'{base}.replay-{horodatage}' + (f'-{n}' if n > 1 else '')
        if not os.path.exists(nom + '.ndjson') and not os.path.exists(nom + '.summary.txt'):
            return nom + '.ndjson'
    raise Erreur('aucun nom de sortie libre à côté de la trace')


def main(argv=None):
    p = argparse.ArgumentParser(prog='etm-powersync-trace-replay',
                                description='Rejoue un fichier de traces de décision vers un optimiseur.')
    p.add_argument('trace', help='fichier de traces (NDJSON, annexe D)')
    p.add_argument('--endpoint', default=ENDPOINT_DEFAUT,
                   help=f'unix:///chemin ou tcp://hôte:port (défaut : {ENDPOINT_DEFAUT})')
    p.add_argument('--token-file', help='fichier contenant le token (jamais en argument)')
    p.add_argument('--out', help='fichier des réponses (défaut : à côté de la trace) ; jamais écrasé')
    p.add_argument('--timeout', type=float, default=30.0, help='délai par requête, en s (défaut 30)')
    p.add_argument('--tolerance-w', type=float, default=0.5,
                   help='écart de puissance toléré, en W (défaut 0,5)')
    p.add_argument('--fail-on-divergence', action='store_true',
                   help='code de sortie 1 si une action diffère du plan enregistré')
    args = p.parse_args(argv)

    try:
        with open(args.trace, encoding='utf-8') as f:
            lignes = [(n, l) for n, l in enumerate(f, 1) if l.strip()]
    except OSError as e:
        print(f'trace illisible : {args.trace} : {e.strerror}', file=sys.stderr)
        return 2

    try:
        sortie = args.out or sortie_libre(args.trace)
    except Erreur as e:
        print(f'erreur : {e}', file=sys.stderr)
        return 2
    resume_chemin = (sortie[:-len('.ndjson')] if sortie.endswith('.ndjson') else sortie) + '.summary.txt'
    try:
        token = lire_token(args.token_file)
        sock = ouvrir_endpoint(args.endpoint, args.timeout)
        client = Client(sock)
        identite = handshake(client, token)
        out = open(sortie, 'x', encoding='utf-8')   # 'x' : jamais d'écrasement
    except FileExistsError:
        print(f'refus d\'écraser {sortie}', file=sys.stderr)
        return 2
    except (Erreur, OSError, ConnectionError) as e:
        print(f'erreur : {e}', file=sys.stderr)
        return 2

    requetes = abstentions = erreurs = illisibles = non_comparables = 0
    compute = []
    divergents = []
    with out:
        for n, brute in lignes:
            try:
                ligne = json.loads(brute)
            except ValueError:
                illisibles += 1
                continue
            requete = ligne.get('request') if isinstance(ligne, dict) else None
            if not isinstance(requete, dict):
                continue        # cycle L2 (`strategy: none`) : rien à rejouer
            requetes += 1
            t0 = time.monotonic()
            try:
                r = client.appeler('GetPlan', requete)
            except (OSError, ConnectionError) as e:
                print(f'erreur à la ligne {n} : {e}', file=sys.stderr)
                return 2
            latence = round((time.monotonic() - t0) * 1000)
            reponse = {'result': r['result']} if 'result' in r else \
                      {'error': {'code': (r.get('error') or {}).get('code')}}
            maintenant = instant(requete.get('timestamp'))
            attendu = actions_courantes(ligne.get('response'), maintenant)
            obtenu = actions_courantes(reponse, maintenant)
            if obtenu is None:
                erreurs += 1
            elif obtenu == []:
                abstentions += 1
            res = reponse.get('result')
            if isinstance(res, dict):
                cms = (res.get('diagnostics') or {}).get('computeMs')
                if isinstance(cms, (int, float)) and not isinstance(cms, bool):
                    compute.append(cms)
            ecarts = None
            if attendu is None or obtenu is None:
                non_comparables += 1
            else:
                ecarts = differences(attendu, obtenu, args.tolerance_w)
                if ecarts:
                    divergents.append((n, ligne.get('t'), ecarts))
            out.write(json.dumps({'line': n, 't': ligne.get('t'), 'response': reponse,
                                  'latencyMs': latence, 'divergent': bool(ecarts),
                                  'differences': ecarts or []},
                                 separators=(',', ':'), ensure_ascii=False) + '\n')
    sock.close()

    resume = [f'trace : {args.trace}',
              f'optimiseur : {str(identite.get("name", "?"))[:64]} ({str(identite.get("protocolVersion"))[:32]})'
              f' sur {args.endpoint}',
              f'réponses : {sortie}',
              f'requêtes rejouées : {requetes}' + (f' ({illisibles} ligne(s) illisible(s) ignorée(s))' if illisibles else ''),
              f'abstentions : {abstentions} ({(100.0 * abstentions / requetes) if requetes else 0:.1f} %)',
              f'erreurs : {erreurs}',
              ('computeMs : médian ' + f'{statistics.median(compute):g}, max {max(compute):g} '
               f'({len(compute)} réponse(s) qui le donnent)') if compute else 'computeMs : non fourni',
              f'cycles non comparables (réponse enregistrée en erreur, ou rejouée en erreur) : {non_comparables}',
              f'cycles dont les actions diffèrent du plan enregistré : {len(divergents)}']
    for n, t, ecarts in divergents:
        resume.append(f'  ligne {n} ({t}) : ' + ' ; '.join(ecarts))
    texte = '\n'.join(resume) + '\n'
    sys.stdout.write(texte)
    try:
        with open(resume_chemin, 'x', encoding='utf-8') as f:
            f.write(texte)
    except OSError as e:
        print(f'résumé non écrit ({resume_chemin}) : {e.strerror}', file=sys.stderr)
    return 1 if (divergents and args.fail_on_divergence) else 0


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