#!/usr/bin/env python3
"""Yasha v200 Colab Runner — single-cell setup, robust against failures."""
import subprocess, sys, os, threading, time, json, random, string, urllib.request, signal

# ── 1. Install deps ──
DEPS = ['torch', 'transformers', 'peft', 'datasets', 'bitsandbytes',
        'accelerate', 'sentencepiece', 'flask', 'huggingface-hub',
        'protobuf', 'tqdm', 'pyarrow', 'requests']
subprocess.run([sys.executable, '-m', 'pip', 'install', '-q'] + DEPS, capture_output=True)

import torch
print(f'Torch CUDA: {torch.cuda.is_available()}')
if torch.cuda.is_available():
    print(f'GPU: {torch.cuda.get_device_name(0)}  VRAM: {torch.cuda.get_device_properties(0).total_memory/1e9:.1f}GB')

# ── 2. HF Token ──
HF_TOKEN = os.environ.get('HF_TOKEN', '')
if not HF_TOKEN:
    try:
        from google.colab import userdata
        HF_TOKEN = userdata.get('HF_TOKEN')
    except: pass
if not HF_TOKEN:
    HF_TOKEN = input('Paste HF_TOKEN (or Enter to skip): ').strip()
os.environ['HF_TOKEN'] = HF_TOKEN or ''
print(f'HF_TOKEN: {"✅ Set" if HF_TOKEN else "❌ Not set"}')

# ── 3. Download engine ──
REPO = 'BeheraBoi/yasha-v200-engine'
FILES = ['arch_v2.py', 'crash_protector.py', 'train_gpu.py']
os.makedirs('/content/yasha-engine', exist_ok=True)
os.makedirs('/content/yasha_v200', exist_ok=True)
for fname in FILES:
    url = f'https://huggingface.co/{REPO}/raw/main/{fname}'
    dest = f'/content/yasha-engine/{fname}' if fname != 'train_gpu.py' else f'/content/yasha_v200/{fname}'
    try:
        urllib.request.urlretrieve(url, dest)
        print(f'✅ {fname}')
    except Exception as e:
        print(f'❌ {fname}: {e}')
open('/content/yasha-engine/__init__.py', 'w').close()

# ── 4. Control Server (Flask) — find free port ──
CONTROL_TOKEN = ''.join(random.choices(string.ascii_letters + string.digits, k=16))
from flask import Flask, request, jsonify, send_file, Response

app = Flask(__name__)
train_process = None
train_log = []
app._train_log = train_log

def check_token():
    t = request.args.get('token', request.headers.get('X-Token', ''))
    if t != CONTROL_TOKEN: return jsonify({'error': 'unauthorized'}), 403
    return None

@app.route('/status')
def status():
    auth = check_token(); 
    if auth: return auth
    alive = train_process and train_process.poll() is None
    return jsonify({'alive': alive, 'pid': train_process.pid if alive else None,
                    'returncode': train_process.returncode if not alive else None,
                    'log_lines': len(train_log), 'last_10': train_log[-20:]})

@app.route('/log')
def get_log():
    auth = check_token();
    if auth: return auth
    n = min(int(request.args.get('n', 50)), 500)
    return Response('\n'.join(train_log[-n:]), mimetype='text/plain')

@app.route('/command', methods=['POST'])
def run_command():
    auth = check_token();
    if auth: return auth
    cmd = request.form.get('cmd', '')
    cwd = request.form.get('cwd', '/content')
    timeout_s = int(request.form.get('timeout', 300))
    try:
        r = subprocess.run(cmd, shell=True, capture_output=True, text=True, cwd=cwd, timeout=timeout_s)
        return jsonify({'stdout': r.stdout, 'stderr': r.stderr, 'returncode': r.returncode})
    except subprocess.TimeoutExpired:
        return jsonify({'stdout': 'TIMEOUT', 'stderr': '', 'returncode': -1})
    except Exception as e:
        return jsonify({'stdout': '', 'stderr': str(e), 'returncode': -1})

@app.route('/upload', methods=['POST'])
def upload_file():
    auth = check_token();
    if auth: return auth
    path = request.form.get('path', '')
    file = request.files.get('file')
    if not file: return jsonify({'error': 'no file'}), 400
    os.makedirs(os.path.dirname(path), exist_ok=True) if '/' in path else None
    file.save(path)
    return jsonify({'path': path, 'size': os.path.getsize(path)})

@app.route('/download')
def download_file():
    auth = check_token();
    if auth: return auth
    path = request.args.get('path', '')
    if not os.path.exists(path): return jsonify({'error': 'not found'}), 404
    return send_file(path, as_attachment=True)

@app.route('/start', methods=['POST'])
def start_training():
    auth = check_token();
    if auth: return auth
    global train_process, train_log
    script = request.form.get('script', '/content/yasha_v200/train_gpu.py')
    if not os.path.exists(script): return jsonify({'error': f'{script} not found'}), 400
    train_log = []
    def runner():
        global train_process
        proc = subprocess.Popen([sys.executable, '-u', script],
            stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
            bufsize=1, text=True, cwd='/content',
            env={**os.environ, 'HF_TOKEN': os.environ.get('HF_TOKEN', '')})
        train_process = proc
        for line in proc.stdout:
            train_log.append(line.rstrip())
            if len(train_log) > 10000: train_log[:] = train_log[-5000:]
            print(line, end='')
        proc.wait()
    threading.Thread(target=runner, daemon=True).start()
    time.sleep(1)
    return jsonify({'started': True, 'script': script})

@app.route('/stop', methods=['POST'])
def stop_training():
    auth = check_token();
    if auth: return auth
    global train_process
    if train_process and train_process.poll() is None:
        train_process.terminate()
        time.sleep(2)
        if train_process.poll() is None: train_process.kill()
        return jsonify({'stopped': True})
    return jsonify({'stopped': False, 'reason': 'not running'})

@app.route('/gpu')
def gpu_status():
    auth = check_token();
    if auth: return auth
    try:
        r = subprocess.run(['nvidia-smi', '--query-gpu=index,name,memory.used,memory.total,temperature.gpu',
                           '--format=csv,noheader'], capture_output=True, text=True, timeout=10)
        return jsonify({'output': r.stdout, 'error': r.stderr})
    except Exception as e: return jsonify({'output': '', 'error': str(e)})

# Find available port
PORT = 8080
for port_try in range(8080, 8100):
    try:
        import socket
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.bind(('0.0.0.0', port_try))
        s.close()
        PORT = port_try
        break
    except: continue

def run_flask(): app.run(host='0.0.0.0', port=PORT, debug=False, use_reloader=False)
flask_thread = threading.Thread(target=run_flask, daemon=True)
flask_thread.start()
time.sleep(1)
print(f'✅ Control server on port {PORT}')

# ── 5. Tunnel via Colab proxy ──
try:
    from google.colab.output import eval_js
    URL = eval_js(f'google.colab.kernel.proxyPort({PORT})')
    print(f'✅ Tunnel: {URL}')
except Exception as e:
    print(f'❌ Colab proxy: {e}')
    URL = f'http://localhost:{PORT}'

print(f'\n🔗 URL:   {URL}')
print(f'🔑 Token: {CONTROL_TOKEN}')
print(f'💬 Send to Aayush\n')

# ── 6. Keep alive ──
while True:
    time.sleep(30)
    try:
        import psutil; psutil.Process()
    except: pass