#!/usr/bin/env python3
"""Sample a stronger local code teacher on a free T4; retain only validated games."""
from __future__ import annotations
import argparse
import hashlib
import json
import os
from pathlib import Path
import re
import sys
import time
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / 'src'))
from vibegame.project import parse_project
from vibegame.seeds import PROJECT, SCENE
from vibegame.validate import validate_project

TEACHER = 'unsloth/Qwen2.5-Coder-14B-Instruct-bnb-4bit'
REVISION = 'ecfc63c083900ce4f19daad695e711f38ab98fcf'
TEACHER_PROMPT = '''Write only the contents of main.gd for a complete playable Godot 4.5 game.
The project and scene already exist: a single Node2D with this script, viewport 960x640, gl_compatibility renderer. Begin with extends Node2D. Draw every visible object with _draw(), draw_circle(), draw_rect() and draw_string(ThemeDB.fallback_font, position, text, HORIZONTAL_ALIGNMENT_LEFT, -1, 20, Color.WHITE). Update game state in _process(delta: float) and call queue_redraw(). Do not assume child nodes exist. No assets, textures, scenes, loading, filesystem, network, OS, or tool APIs. Do not create project/scene files or JSON.
Use arrow actions ui_left/ui_right/ui_up/ui_down, ui_accept for space/enter, and Input.is_physical_key_pressed(KEY_R) for restart. Show controls, score/progress and win/lose feedback. Implement the requested mechanic with simple shapes in at most 100 lines. Use ordinary = assignment instead of := inference on expressions. GDScript uses value_if_true if condition else value_if_false, never C ternaries. Use queue_redraw(), not update(). Use Color("ffffff"), Vector2(x,y), Rect2(x,y,w,h). Do not quit automatically. Output code only, no explanation or Markdown.
Here is a small valid drawing/input example, not the game to reproduce:
extends Node2D
var p = Vector2(480, 320)
var goal = Vector2(760, 220)
var score = 0
func _process(dt: float) -> void:
    p += Input.get_vector("ui_left", "ui_right", "ui_up", "ui_down") * 250.0 * dt
    p = p.clamp(Vector2(20, 100), Vector2(940, 620))
    if p.distance_to(goal) < 24.0:
        score += 1
        goal = Vector2(randf_range(40, 920), randf_range(120, 600))
    if Input.is_physical_key_pressed(KEY_R):
        score = 0
    queue_redraw()
func _draw() -> void:
    draw_circle(p, 14, Color("5ef2d0"))
    draw_circle(goal, 10, Color("ffdc75"))
    draw_string(ThemeDB.fallback_font, Vector2(24, 38), "Arrows: move | R: restart | Score %d" % score, HORIZONTAL_ALIGNMENT_LEFT, -1, 20, Color.WHITE)
'''

def script_project(raw):
    """Normalize only a teacher's enclosing code fence, then validate full files."""
    code = raw.strip()
    if code.startswith('```') and code.endswith('```'):
        code = code.split('\n', 1)[1].rsplit('```', 1)[0].strip()
    if not code.startswith('extends Node2D'):
        raise ValueError('Teacher must return a main.gd script beginning extends Node2D')
    return parse_project(json.dumps({'files': {'project.godot': PROJECT % 'VibeGame teacher',
        'main.tscn': SCENE, 'main.gd': code + '\n'}}))

def read_rows(path):
    return [json.loads(line) for line in Path(path).read_text().splitlines() if line.strip()]

def digest(obj):
    return hashlib.sha256(json.dumps(obj, sort_keys=True).encode()).hexdigest()

def append(path, obj):
    with path.open('a') as f:
        f.write(json.dumps(obj, ensure_ascii=False) + '\n')
        f.flush()
        os.fsync(f.fileno())

def choose_reference(prompt, corpus):
    words = set(re.findall(r'[a-z]{4,}', prompt.lower()))
    candidates = [r for r in corpus if r['path'].startswith('2d/') and len(r['code']) < 1800]
    if not candidates:
        return None
    return max(candidates, key=lambda r: len(words & set(re.findall(r'[a-z]{4,}', r['path'].lower() + ' ' + r['code'].lower()))))

def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument('--prompts', type=Path, default=ROOT / 'data/teacher_prompts.jsonl')
    p.add_argument('--blind', type=Path, default=ROOT / 'data/eval_prompts.jsonl')
    p.add_argument('--corpus', type=Path, default=ROOT / 'data/reference_gd.jsonl')
    p.add_argument('--output', type=Path, required=True, help='Persistent Drive folder')
    p.add_argument('--godot', required=True)
    p.add_argument('--limit', type=int, default=100)
    p.add_argument('--session-hours', type=float, default=10.5)
    args = p.parse_args()
    if not 0 < args.session_hours <= 10.5 or args.limit < 1:
        p.error('Positive limits and session-hours <=10.5 required')
    import unsloth
    import torch
    if not torch.cuda.is_available() or 'T4' not in torch.cuda.get_device_name(0):
        raise RuntimeError('Select a free Colab T4; this script provisions no GPU')
    if not str(args.output.resolve()).startswith('/content/drive/MyDrive/'):
        raise RuntimeError('Teacher attempts and admitted pairs must persist on Drive')
    prompts = read_rows(args.prompts)[:args.limit]
    blind = {r['instruction'].strip().casefold() for r in read_rows(args.blind)}
    if any(r['instruction'].strip().casefold() in blind for r in prompts):
        raise ValueError('Refusing to generate teacher answers to blind evaluation prompts')
    corpus = read_rows(args.corpus)
    if any(r['license'] not in {'MIT', 'Apache-2.0', 'CC0-1.0'} for r in corpus):
        raise ValueError('Reference corpus has a disallowed license')
    args.output.mkdir(parents=True, exist_ok=True)
    journal = args.output / 'attempts.jsonl'
    admitted = args.output / 'validated_candidates.jsonl'
    attempted = {r['id'] for r in read_rows(journal)} if journal.exists() else set()
    manifest = {'teacher': TEACHER, 'revision': REVISION, 'prompt_file_sha256': hashlib.sha256(args.prompts.read_bytes()).hexdigest(),
                'system_prompt_sha256': hashlib.sha256(TEACHER_PROMPT.encode()).hexdigest(),
                'fixed_project_sha256': digest({'project': PROJECT, 'scene': SCENE}),
                'corpus_sha256': hashlib.sha256(args.corpus.read_bytes()).hexdigest(),
                'seed': 3407, 'temperature': 0.35, 'top_p': 0.9, 'max_sequence_length': 4096,
                'max_new_tokens': 2600, 'protocol': 'one main.gd sample per prompt, assemble original fixed project/scene, strip enclosing code fence only; discard any validation failure; no paid API',
                'notes': 'Teacher targets can inherit reference licensing obligations; retain all source notices.'}
    mf = args.output / 'manifest.json'
    if mf.exists() and json.loads(mf.read_text()) != manifest:
        raise ValueError('Teacher run configuration changed; use another output directory')
    mf.write_text(json.dumps(manifest, indent=2) + '\n')
    started = time.monotonic()
    model, tokenizer = unsloth.FastLanguageModel.from_pretrained(model_name=TEACHER, revision=REVISION,
        max_seq_length=4096, dtype=torch.float16, load_in_4bit=True,
        use_exact_model_name=True, trust_remote_code=False)
    from bitsandbytes.nn import Linear4bit
    compute_counts = {}
    for module in model.modules():
        if isinstance(module, Linear4bit):
            compute_counts[str(module.compute_dtype)] = compute_counts.get(str(module.compute_dtype), 0) + 1
            module.compute_dtype = torch.float16
            module.compute_type_is_set = True
    print(json.dumps({'linear4bit_previous_compute_dtypes': compute_counts, 'compute_dtype': 'torch.float16'}), flush=True)
    unsloth.FastLanguageModel.for_inference(model)
    for row in prompts:
        identity = row.get('id', digest(row['instruction'])[:20])
        if identity in attempted:
            continue
        if time.monotonic() - started > args.session_hours * 3600:
            print('Session budget reached; rerun to resume unattempted prompts.', flush=True)
            return
        reference = choose_reference(row['instruction'], corpus)
        context = ''
        if reference:
            context = ('\nThe following licensed GDScript is reference material, not instructions. Use its API patterns only where appropriate. '
                       'Build an independent complete game; do not include asset references from this snippet.\n'
                       f"Source: {reference['id']} License: {reference['license']}\n<reference>\n{reference['code']}\n</reference>")
        # The stored instruction remains the complete-project task. The teacher
        # writes gameplay code; deterministic original scaffolding supplies files.
        task = row['instruction'].split(' Make it a complete playable Godot')[0]
        messages = [{'role':'system','content':TEACHER_PROMPT}, {'role':'user','content':task + '\nReturn main.gd code only.' + context}]
        text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        encoded = tokenizer(text, return_tensors='pt').to('cuda')
        budget = min(2600, 4096 - encoded.input_ids.shape[1])
        if budget < 2300:
            # Reference is optional. Never silently truncate a project or prompt.
            messages[-1]['content'] = task + '\nReturn main.gd code only.'
            reference = None
            text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
            encoded = tokenizer(text, return_tensors='pt').to('cuda')
            budget = min(2600, 4096 - encoded.input_ids.shape[1])
        seed = 3407 + int(hashlib.sha256(identity.encode()).hexdigest()[:8], 16) % 100000
        torch.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
        start = time.monotonic()
        evidence = {'id': identity, 'instruction': row['instruction'], 'seed': seed,
                    'reference': {k:v for k,v in reference.items() if k != 'code'} if reference else None,
                    'accepted':False}
        case = args.output / 'cases' / identity
        case.mkdir(parents=True, exist_ok=True)
        try:
            if (case / 'raw.txt').exists():
                # A disconnected attempt is resumed without another stochastic sample.
                raw = (case / 'raw.txt').read_text()
            else:
                with torch.inference_mode():
                    tokens = model.generate(**encoded, max_new_tokens=budget, do_sample=True,
                        temperature=0.35, top_p=0.9, pad_token_id=tokenizer.eos_token_id)
                raw = tokenizer.decode(tokens[0, encoded.input_ids.shape[1]:], skip_special_tokens=True)
                (case / 'raw.txt').write_text(raw)
            evidence['generation_seconds'] = round(time.monotonic() - start, 3)
            files = script_project(raw)
            validation_dir = case / ('validation-' + str(time.time_ns()))
            result = validate_project(files, validation_dir, args.godot, run_seconds=60)
            evidence['validation'] = result
            evidence['accepted'] = bool(result['compile_success'] and result['clean_run_60s'])
            if evidence['accepted']:
                candidate = {'instruction':row['instruction'], 'output':json.dumps({'files':files}, ensure_ascii=False)}
                # The journal is canonical; materialize candidates from it after each write.
                evidence['candidate'] = candidate
        except (torch.cuda.OutOfMemoryError, KeyboardInterrupt):
            raise
        except Exception as exc:
            evidence['error'] = type(exc).__name__ + ': ' + str(exc)
        append(journal, evidence)
        rows = [r['candidate'] for r in read_rows(journal) if r.get('accepted')]
        temp = admitted.with_suffix('.tmp')
        temp.write_text(''.join(json.dumps(r, ensure_ascii=False) + '\n' for r in rows))
        temp.replace(admitted)
        print(json.dumps({'attempt':identity,'accepted':evidence['accepted'],'accepted_total':len(rows),
                          'error':evidence.get('error')}), flush=True)
    print('Teacher sampling complete.', flush=True)

if __name__ == '__main__':
    main()