#!/usr/bin/env python3
"""Persistent server for OUR self-starter — cold-start romantic piano from silence.
Backend for the poly_start MCP tool.

CONDITIONED on composer + key + (now) ATTRIBUTES. Loads ckpt_attr.pt + attr_meta.json and
prefixes generation with [instr, composer, key, <S>, <9 attribute tokens>]; the agent maps a
brief's language ("a shimmering wave that descends in a dark stormy mood") to attribute bins.
Prefix/attribute tokens stay in context but are stripped before detok and masked from sampling.
Falls back to ckpt_long300 (composer+key only) then best_big.pt if the attr model is absent.
  tmp/melody/aria-repo/.venv/bin/python tmp/melody/poly_server.py

POST /coldstart { composer?, key?, attributes?{attr:bin}, tokens?, temp?, topp?, variations?, cfg? }
  -> { conditioned, attrs, variations: [ { notes:[{pitch,start,end,velocity}], composer, key, attributes } ] }
GET /health
"""
import os, sys, json, io, random
from http.server import BaseHTTPRequestHandler, HTTPServer
import torch
import pretty_midi

ROOT = '/Users/alhill/projects/music/tmp/melody'
sys.path.insert(0, ROOT); sys.path.insert(0, f'{ROOT}/aria-repo')
from model_qwen import GPTQwen
from ariautils.tokenizer import AbsTokenizer

CANDS = [(f'{ROOT}/ckpt_attr.pt', f'{ROOT}/attr_meta.json'),          # attribute-conditioned (preferred)
         (f'{ROOT}/ckpt_long300.pt', f'{ROOT}/cond_meta_long.json'),  # composer+key only
         (f'{ROOT}/best_big.pt', None)]                               # unconditional
CKPT, METAF = os.environ.get('POLY_CKPT'), os.environ.get('POLY_META')
if not CKPT:
    for c, m in CANDS:
        if os.path.exists(c): CKPT, METAF = c, m; break

TOK = AbsTokenizer()
CK = torch.load(CKPT, map_location='cpu'); A = CK['arch']
META = json.load(open(METAF)) if METAF else None
VOCAB = META['vocab'] if META else CK['meta']['vocab']
DEV = torch.device('mps' if torch.backends.mps.is_available() else 'cpu')
MODEL = GPTQwen(vocab=VOCAB, d=A['d'], nh=A['nh'], nkv=A['nkv'], nl=A['nl'], ctx=A['ctx']).to(DEV)
MODEL.load_state_dict(CK['model']); MODEL.eval()

COND = bool(META and 'composers' in META)
if COND:
    INSTR, BOS, EOS = META['instr_prefix_id'], META['bos'], META['eos']
    COMPOSERS, KEYS = META['composers'], META['keys']
    STRIP = set(COMPOSERS.values()) | set(KEYS.values()) | {INSTR, BOS, EOS, META.get('pad_id', -1)}
else:
    INSTR = TOK.encode([('prefix', 'instrument', 'piano')])[0]
    BOS, EOS = TOK.tok_to_id[TOK.bos_tok], TOK.tok_to_id[TOK.eos_tok]
    COMPOSERS, KEYS, STRIP = {}, {}, {INSTR, BOS, EOS}

ATTR = META.get('attr_order') if META else None                      # attribute schema present?
if ATTR:
    ATTR_TOK, ATTR_ANY, ATTR_SCHEMA = META['attr_tokens'], META['attr_any'], META['attr_schema']
    for at in ATTR:
        STRIP |= set(ATTR_TOK[at].values()); STRIP.add(ATTR_ANY[at])
SMASK = torch.zeros(VOCAB)                                            # never GENERATE a prefix/attr token (but allow EOS)
for i in (STRIP - {EOS}):
    if 0 <= i < VOCAB: SMASK[i] = float('-inf')

print(f"[poly] ready ({os.path.basename(CKPT)} step {CK.get('step')} val {CK.get('val'):.3f}, "
      f"{sum(p.numel() for p in MODEL.parameters())/1e6:.0f}M params, "
      f"{len(COMPOSERS)} composers/{len(KEYS)} keys" + (f", {len(ATTR)} ATTRS" if ATTR else "") + ")", flush=True)


def seed_for(composer, key, attributes=None):
    if not COND:
        return [INSTR, BOS], ('', '', {})
    comp = composer if composer in COMPOSERS else 'unknown'
    k = key if key in KEYS else random.choice(list(KEYS))
    seed = [INSTR, COMPOSERS[comp], KEYS[k], BOS]
    used = {}
    if ATTR:                                                         # append the attribute block (unspecified -> 'any')
        attributes = attributes or {}
        for at in ATTR:
            b = attributes.get(at)
            if b in ATTR_TOK[at]:
                seed.append(ATTR_TOK[at][b]); used[at] = b
            else:
                seed.append(ATTR_ANY[at])
    return seed, (comp, k, used)


def sample(logits, temp, topp):
    p = torch.softmax((logits.to('cpu').float() + SMASK) / temp, dim=-1)   # mask prefix/attr; sample on CPU
    sp, si = torch.sort(p, descending=True)
    sp = sp * ((torch.cumsum(sp, 0) - sp) < topp)
    return si[torch.multinomial(sp / sp.sum(), 1)].item()


def gen(seed_ids, n_tokens, temp, topp):
    toks = list(seed_ids); note_ids = []
    logits, cache = MODEL.infer(torch.tensor([toks], device=DEV), cache=None, pos=0)
    for _ in range(n_tokens):
        nxt = sample(logits[0], temp, topp)
        if nxt == EOS or len(toks) >= MODEL.ctx - 1:
            break
        toks.append(nxt)
        if nxt not in STRIP:
            note_ids.append(nxt)
        logits, cache = MODEL.infer(torch.tensor([[nxt]], device=DEV), cache=cache, pos=len(toks) - 1)
    return note_ids


def to_notes(note_ids):
    md = TOK.detokenize(TOK.decode([INSTR, BOS] + note_ids))
    buf = io.BytesIO(); md.to_midi().save(file=buf)
    pm = pretty_midi.PrettyMIDI(io.BytesIO(buf.getvalue()))
    notes = sorted((n.start, n.end, n.pitch, n.velocity) for inst in pm.instruments for n in inst.notes)
    return [{'pitch': p, 'start': round(s, 3), 'end': round(e, 3), 'velocity': v} for (s, e, p, v) in notes]


def coldstart(composer, key, attributes, tokens, temp, topp, variations):
    out = []
    for _ in range(variations):
        seed, (comp, k, used) = seed_for(composer, key, attributes)
        nids = gen(seed, int(tokens), temp, topp)
        out.append({'notes': to_notes(nids), 'composer': comp, 'key': k, 'attributes': used})
    return out


class H(BaseHTTPRequestHandler):
    def _send(self, code, obj):
        b = json.dumps(obj).encode()
        self.send_response(code); self.send_header('Content-Type', 'application/json')
        self.send_header('Content-Length', str(len(b))); self.end_headers(); self.wfile.write(b)

    def do_GET(self):
        if self.path == '/health':
            self._send(200, {'ok': True, 'ckpt': os.path.basename(CKPT), 'step': CK.get('step'), 'val': CK.get('val'),
                             'conditioned': COND, 'composers': sorted(COMPOSERS), 'keys': sorted(KEYS),
                             'attributes': ATTR_SCHEMA if ATTR else None})
        else:
            self._send(404, {'error': 'not found'})

    def do_POST(self):
        if self.path != '/coldstart':
            return self._send(404, {'error': 'not found'})
        try:
            req = json.loads(self.rfile.read(int(self.headers.get('Content-Length', 0))))
            self._send(200, {'conditioned': COND, 'attrs': bool(ATTR), 'variations': coldstart(
                req.get('composer', 'unknown'), req.get('key'), req.get('attributes'),
                int(req.get('tokens', 1024)),                             # CONFIGURABLE (honored from request)
                float(req.get('temp', 1.0)), req.get('topp', 0.96), int(req.get('variations', 1)))})
        except Exception as e:
            import traceback; traceback.print_exc(); self._send(500, {'error': str(e)})

    def log_message(self, *a): pass


if __name__ == '__main__':
    port = int(os.environ.get('POLY_PORT', '8793'))
    print(f'[poly] serving on 127.0.0.1:{port}', flush=True)
    HTTPServer(('127.0.0.1', port), H).serve_forever()