from __future__ import annotations

import csv
import datetime
import io
import json
import os
import sys
import urllib.request

DEFAULT_BASE = 'https://trends.polyhelper.ai/api/v1/'
LAST = {}
SKIP_RULES = ('f12',)

def read(base, name):
    if base.startswith(('http://', 'https://')):
        req = urllib.request.Request(base.rstrip('/') + '/' + name, headers={'User-Agent': 'poly-trends-replay/1'})
        with urllib.request.urlopen(req, timeout=60) as r:
            return r.read().decode('utf-8')
    with open(os.path.join(base, name), encoding='utf-8') as f:
        return f.read()

def rows_of(text):
    return list(csv.DictReader(io.StringIO(text)))

def num(x, d=0.0):
    try:
        return float(x)
    except (TypeError, ValueError):
        return d

def px_table(prices):
    out = {}
    for s in prices.get('sessions') or []:
        out[s['d']] = {k: tuple(v) for k, v in (s.get('px') or {}).items()}
    f12 = prices.get('f12') or {}
    funds = f12.get('funds') or {}
    for d, m in (f12.get('px') or {}).items():
        for code, v in (m or {}).items():
            pair = tuple(v) if isinstance(v, (list, tuple)) else (v.get('exec'), v.get('close')) if isinstance(v, dict) else None
            if pair:
                out.setdefault(d, {})[funds.get(code, code)] = pair
                out[d].setdefault(code, pair)
    return out

def days_between(a, b):
    return (datetime.date.fromisoformat(b) - datetime.date.fromisoformat(a)).days

def replay_account(acct, sessions, trades, px, rates, corp, pay_lag, start_cash, pro=None):
    cash, units, last = start_cash, {}, {}
    fees, due, done_ev, held_log = 0.0, [], set(), []
    out, prev, n_trades = [], None, 0
    for d in sessions:
        problems = []
        if prev is not None and pro:
            r_, n_ = rates.get(d, 0.0), days_between(prev, d)
            fin = cash * r_ * n_ / 360 if cash >= 0 else cash * (r_ + pro['spread']) * n_ / 360
            shorts = sum(-u * (last.get(k) or 0.0) for k, u in units.items() if u < 0)
            cash += fin - shorts * pro['borrow'] * n_ / 360
        elif prev is not None and cash > 0:
            cash += cash * rates.get(d, 0.0) * days_between(prev, d) / 360
        if prev is not None:
            def at_ex(code, ex):
                rec = [u for dd, u in held_log if dd < ex]
                return (rec[-1] if rec else units).get(code, 0.0)
            for code in sorted(set(units) | {c for _, u in held_log for c in u}):
                ev = corp.get(code) or {}
                for day, ratio in ev.get('split') or []:
                    key = f'{code}:{day}'
                    if prev < day <= d and key not in done_ev and ratio > 0:
                        if code in units:
                            units[code] += at_ex(code, day) * (ratio - 1)
                        for dd, u in held_log:
                            if code in u and dd < day:
                                u[code] *= ratio
                        if code in last:
                            last[code] /= ratio
                        done_ev.add(key)
                for day, amt in ev.get('div') or []:
                    key = f'{code}:div:{day}'
                    if prev < day <= d and key not in done_ev and at_ex(code, day):
                        due.append({'ex': day, 'pay': (datetime.date.fromisoformat(day) + datetime.timedelta(days=pay_lag)).isoformat(),
                                    'usd': round(at_ex(code, day) * amt, 2)})
                        done_ev.add(key)
            paid = [x for x in due if x['pay'] <= d]
            cash += sum(x['usd'] for x in paid)
            due = [x for x in due if x['pay'] > d]
        for t in trades.get((d, acct), []):
            inst, side = t['instrument'], t['side']
            price, amount, cost = num(t['price_usd']), num(t['amount_usd']), num(t['cost_usd'])
            sysd = t.get('system') in ('margin', 'bust')
            exec_ = (px.get(d) or {}).get(inst, (None, None))[1 if sysd else 0]
            if exec_ is None:
                problems.append(f'no execution price for {inst}')
            elif abs(exec_ - price) > 1e-6 * max(1.0, price):
                problems.append(f'{inst}: trade price {price} is not the session {"closing" if sysd else "execution"} price {exec_}')
            if not price:
                problems.append(f'{inst}: no price in the trade')
                continue
            du = amount / price
            if side == 'buy':
                units[inst] = units.get(inst, 0.0) + du
                cash -= amount
            else:
                held = units.get(inst, 0.0)
                units[inst] = 0.0 if abs(held - num(t['units'])) <= 1e-4 * max(1.0, held) else held - du
                cash += amount
            cash -= cost
            fees += cost
            last[inst] = price
            n_trades += 1
        units = {k: u for k, u in units.items() if abs(u) > 1e-12}
        for inst, pair in (px.get(d) or {}).items():
            if pair[1] is not None and (inst in units or inst in last):
                last[inst] = pair[1]
        v = cash + sum(x['usd'] for x in due) + sum(u * (last.get(k) or 0.0) for k, u in units.items())
        held_log = (held_log + [(d, dict(units))])[-20:]
        out.append({'session': d, 'value': v, 'cash_share': cash / v if v else 1.0, 'fees': fees, 'n_trades': n_trades, 'problems': problems,
                    'weights': {k: u * (last.get(k) or 0.0) / v for k, u in units.items()} if v else {}})
        prev = d
    return out

def check(base, tol=0.01, lines=None):
    P = (lambda s='': lines.append(s)) if lines is not None else print
    acc = rows_of(read(base, 'arena-accounts.csv'))
    trd = rows_of(read(base, 'arena-trades.csv'))
    prices = json.loads(read(base, 'arena-prices.json'))
    px = px_table(prices)
    rates = {s['d']: num(s.get('rate')) for s in prices.get('sessions') or []}
    corp_all = prices.get('corp') or {}
    start = num(prices.get('start'), 100000.0)
    pay_lag = int(prices.get('pay_lag') or 10)
    pk = (prices.get('f12') or {}).get('pro') or {}
    pro_k = {'borrow': num(pk.get('borrow'), 0.01), 'spread': num(pk.get('spread'), 0.01)}
    trades = {}
    for t in trd:
        trades.setdefault((t['session'], t['account']), []).append(t)
    by = {}
    for r in acc:
        by.setdefault(r['account'], []).append(r)
    all_sessions = sorted({r['session'] for r in acc})
    P(f'replay: {len(acc)} published rows, {len(by)} accounts, {len(all_sessions)} sessions ({all_sessions[0] if all_sessions else "-"} .. '
      f'{all_sessions[-1] if all_sessions else "-"}), {len(trd)} trades, rules {prices.get("rules")}')
    cash_book = {r['session']: num(r['value_usd']) for r in by.get('_cash', [])}
    stats = {'exact': 0, 'rounding': 0, 'diff': 0, 'skipped': 0, 'rows': 0}
    worst = (0.0, '')
    bad = []
    for acct in sorted(by):
        rows = sorted(by[acct], key=lambda r: r['session'])
        if all(r['rules'].startswith(SKIP_RULES) for r in rows) and not px.get(rows[0]['session'], {}).get('SPY'):
            stats['skipped'] += len(rows)
            continue
        sess = [r['session'] for r in rows]
        prev_all = [s for s in all_sessions if s < sess[0]]
        starts = [start] + ([cash_book[prev_all[-1]]] if prev_all and prev_all[-1] in cash_book else [])
        if prev_all:
            starts = starts[::-1]
        best = None
        for s0 in starts:
            rep = replay_account(acct, sess, trades, px, rates, corp_all, pay_lag, s0, pro_k if rows[0]['rules'].startswith('f12p') else None)
            ndiff = sum(1 for r, e in zip(rows, rep) if abs(num(r['value_usd']) - e['value']) > 0.01 + 0.005 * e['n_trades'])
            if best is None or ndiff < best[0]:
                best = (ndiff, rep, s0)
            if ndiff == 0:
                break
        _, rep, s0 = best
        for r, e in zip(rows, rep):
            stats['rows'] += 1
            pub = num(r['value_usd'])
            dv = e['value'] - pub
            msgs = list(e['problems'])
            if abs(dv) > worst[0]:
                worst = (abs(dv), f'{acct} {e["session"]}')
            bound = tol + 0.005 * e['n_trades']
            if abs(dv) <= tol:
                kind = 'exact'
            elif abs(dv) <= bound:
                kind = 'rounding'
            else:
                kind = 'diff'
                msgs.append(f'value: published {pub:.2f}, recounted {e["value"]:.2f} (difference {dv:+.2f})')
            if r['cash_share'] != '' and abs(num(r['cash_share']) - e['cash_share']) > 0.0001:
                kind = 'diff'
                msgs.append(f'cash_share: published {r["cash_share"]}, recounted {e["cash_share"]:.4f}')
            if r['fees_usd'] != '' and abs(num(r['fees_usd']) - e['fees']) > bound + 0.005:
                kind = 'diff'
                msgs.append(f'fees: published {r["fees_usd"]}, recounted {e["fees"]:.2f}')
            wj = r.get('weights_json') or ''
            if wj:
                try:
                    pw = json.loads(wj)
                    for k in sorted(set(pw) | set(e['weights'])):
                        if abs(pw.get(k, 0.0) - e['weights'].get(k, 0.0)) > 0.0002:
                            kind = 'diff'
                            msgs.append(f'weight {k}: published {pw.get(k, 0.0)}, recounted {e["weights"].get(k, 0.0):.4f}')
                except ValueError:
                    msgs.append('weights_json unreadable')
            if msgs and kind != 'diff':
                kind = 'diff'
            stats[kind] += 1
            if kind == 'diff':
                bad.append((acct, e['session'], msgs, s0))
    P(f'exact to the cent (|d| <= {tol}): {stats["exact"]}; within the rounding of the published trades: {stats["rounding"]}; '
      f'DIFFERENT: {stats["diff"]}; not recounted (other price sets): {stats["skipped"]}')
    P(f'largest difference: {worst[0]:.4f} USD ({worst[1] or "-"})')
    for acct, d, msgs, s0 in bad[:200]:
        P(f'DIFF {acct} {d} (start {s0:.2f}): ' + '; '.join(msgs))
    if len(bad) > 200:
        P(f'... and {len(bad) - 200} more')
    P('result: ' + ('the published scores are reproduced' if not stats['diff'] else 'DIFFERENCES FOUND'))
    LAST.clear()
    LAST.update({'ran_on': datetime.date.today().isoformat(), 'sessions': len(all_sessions), 'last_session': all_sessions[-1] if all_sessions else None,
                 'accounts': len(by), 'rows': stats['rows'], 'trades': len(trd), 'exact': stats['exact'], 'rounding': stats['rounding'],
                 'diff': stats['diff'], 'skipped': stats['skipped'], 'max_diff_usd': round(worst[0], 4), 'ok': not stats['diff']})
    if not base.startswith(('http://', 'https://')):
        try:
            LAST['files_at'] = datetime.datetime.fromtimestamp(os.path.getmtime(os.path.join(base, 'arena-accounts.csv')), datetime.timezone.utc).strftime('%Y-%m-%d %H:%M')
        except OSError:
            pass
    return (1 if stats['diff'] else 0), ''

def main(argv):
    tol, args, report = 0.01, [], None
    it = iter(argv)
    for a in it:
        if a == '--tol':
            tol = float(next(it))
        elif a == '--report':
            report = next(it)
        else:
            args.append(a)
    base = args[0] if args else DEFAULT_BASE
    try:
        code, _ = check(base, tol)
    except (OSError, ValueError, KeyError) as e:
        print(f'replay: cannot read the open files at {base}: {type(e).__name__}: {e}', file=sys.stderr)
        return 2
    if report and LAST:
        with open(report, 'w', encoding='utf-8') as f:
            json.dump(LAST, f, ensure_ascii=False, indent=1)
    return code

if __name__ == '__main__':
    sys.exit(main(sys.argv[1:]))