from fastapi import FastAPI, BackgroundTasks
import subprocess
import threading
import os
from fastapi.responses import PlainTextResponse

app = FastAPI()

# Global state to track training
training_status = {
    "status": "idle",
    "log": ""
}

def run_training():
    global training_status
    training_status["status"] = "running"
    training_status["log"] = "Started training...\n"
    
    # Run the trainer script and capture output
    process = subprocess.Popen(
        ["python", "-m", "trainer.train", "--steps", "300"],
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        bufsize=1,
        universal_newlines=True
    )
    
    for line in process.stdout:
        training_status["log"] += line
        
    process.wait()
    training_status["status"] = f"finished with code {process.returncode}"
    training_status["log"] += f"\nTraining finished with exit code: {process.returncode}\n"

@app.get("/")
def read_root():
    return {
        "message": "CodeForge GRPO Training Node Active",
        "status": training_status["status"],
        "endpoints": {
            "/start": "Start training (runs in background)",
            "/logs": "View live training logs"
        }
    }

@app.post("/start")
def start_training(background_tasks: BackgroundTasks):
    if training_status["status"] == "running":
        return {"message": "Training is already running!"}
        
    # Run in a separate thread so it doesn't block the API
    thread = threading.Thread(target=run_training)
    thread.start()
    return {"message": "Training started! Go to /logs to monitor."}

@app.get("/logs", response_class=PlainTextResponse)
def get_logs():
    return training_status["log"]