#!/usr/bin/env python3
"""Live teaching server — the main agent personally drives the env step-by-step.

IPC (file-based, /content/teach/):
  cmd.txt  : 'START <game_id>' | 'RESET' | 'PING' | '<ACTIONn> [x y]'
  obs.json : latest observation snapshot (grid hex rows + change list)
  err.json : last error (if any)
Trace: collect_selfplay-compatible JSONL (grids_before/after, levels, hash,
state_text) so data/build_dataset.py consumes teaching data unchanged.

NOTE: grids[0] is the CURRENT frame (n_frames history after it) — matches
build_candidates/collect_selfplay usage.
"""
import json
import os
import sys
import time
import traceback

sys.path.insert(0, '/content/qwenjev')
sys.path.insert(0, '/content/qwenjev/train')
os.chdir('/content/qwenjev')

from data.collect_selfplay import make_env, meta_from, grids_from
from data.mock_env import step_env
from qwenjev.state_codec import canonical_state, grid_to_text, frame_hash

TEACH = '/content/teach'
CMD = TEACH + '/cmd.txt'
OBS = TEACH + '/obs.json'
ERR = TEACH + '/err.json'
os.makedirs(TEACH, exist_ok=True)
for f in (CMD, OBS, ERR):
    if os.path.exists(f):
        os.remove(f)


def render(grids):
    g = grids[0]
    return [''.join('%x' % (int(v) & 15) for v in row) for row in g]


def changes(pre_cur, cur, limit=48):
    out = []
    if pre_cur is None:
        return out
    R = max(len(pre_cur), len(cur))
    for r in range(min(R, 64)):
        pr = pre_cur[r] if r < len(pre_cur) else []
        cr = cur[r] if r < len(cur) else []
        C = max(len(pr), len(cr))
        for c in range(min(C, 64)):
            a = pr[c] if c < len(pr) else None
            b = cr[c] if c < len(cr) else None
            if a != b:
                sa = ('%x' % (int(a) & 15)) if a is not None else '.'
                sb = ('%x' % (int(b) & 15)) if b is not None else '.'
                out.append('%d,%d:%s>%s' % (r, c, sa, sb))
                if len(out) >= limit:
                    return out
    return out


def write_obs(env, gid, step_no, prev_grids, terminal=None, extra=None):
    grids = grids_from(env)
    meta = meta_from(env)
    s = {
        'ok': True, 'game': gid, 'step': step_no, 'state': meta['state'],
        'levels': meta['levels_completed'], 'win_levels': meta['win_levels'],
        'available': meta['available_actions'],
        'grid': render(grids),
        'shape': [len(grids[0]), len(grids[0][0])],
        'n_frames': len(grids),
        'changes': changes(prev_grids[0] if prev_grids else None, grids[0]),
    }
    if terminal:
        s['terminal'] = terminal
    if extra:
        s.update(extra)
    tmp = OBS + '.tmp'
    open(tmp, 'w').write(json.dumps(s, separators=(',', ':')))
    os.replace(tmp, OBS)


env = None
game_id = None
step_no = 0
trace_path = None
prev_grids = None

while True:
    try:
        if os.path.exists(ERR):
            os.remove(ERR)
        if os.path.exists(CMD):
            line = open(CMD).read().strip()
            os.remove(CMD)
            if line:
                parts = line.split()
                if parts[0] == 'PING':
                    if env is not None:
                        write_obs(env, game_id, step_no, prev_grids)
                    else:
                        open(OBS, 'w').write(json.dumps({'ok': True, 'idle': True}))
                elif parts[0] == 'START':
                    game_id = parts[1]
                    trace_path = TEACH + '/trace_teach_%s.jsonl' % game_id
                    env = make_env(game_id, None)
                    env.reset()
                    step_no = 0
                    prev_grids = None
                    write_obs(env, game_id, step_no, prev_grids,
                              extra={'trace': trace_path, 'episode_started': True})
                elif parts[0] == 'RESET':
                    env.reset()
                    prev_grids = None
                    write_obs(env, game_id, step_no, prev_grids, extra={'reset': True})
                else:
                    action = parts[0]
                    x = int(parts[1]) if len(parts) > 1 and parts[1] != '-' else None
                    y = int(parts[2]) if len(parts) > 2 and parts[2] != '-' else None
                    grids = grids_from(env)
                    meta = meta_from(env)
                    step_env(env, action, x, y)
                    step_no += 1
                    grids_after = grids_from(env)
                    meta_after = meta_from(env)
                    record = {
                        'game_id': game_id, 'step': step_no - 1,
                        'action': action, 'x': x, 'y': y,
                        'grids_before': grids, 'grids_after': grids_after,
                        'levels_completed_before': meta['levels_completed'],
                        'levels_completed_after': meta_after['levels_completed'],
                        'available_actions': list(meta['available_actions']),
                        'frame_hash': frame_hash(grids[0]),
                        'state_text': canonical_state(grid_to_text(grids[0])),
                        'ts': time.time(),
                    }
                    open(trace_path, 'a').write(json.dumps(record) + '\n')
                    term = meta_after['state'] if meta_after['state'] in ('WIN', 'GAME_OVER') else None
                    write_obs(env, game_id, step_no, grids, terminal=term)
                    if term == 'GAME_OVER':
                        env.reset()
                        prev_grids = None
                    else:
                        prev_grids = grids_after
        time.sleep(0.2)
    except Exception as e:
        open(ERR, 'w').write(json.dumps(
            {'error': '%s: %s' % (type(e).__name__, e),
             'tb': traceback.format_exc()[-1500:]}))
        time.sleep(0.5)