import os
from fastapi import FastAPI
from lightning_sdk import Studio
import logging
from io import StringIO

app = FastAPI()

log_stream = StringIO()
logging.basicConfig(stream=log_stream, level=logging.INFO)
    
@app.post("/start-studio/")
def start_studio(userid: str, apikey: str, name: str, teamspace: str, user: str):
    try:
        os.environ["LIGHTNING_USER_ID"] = userid
        os.environ["LIGHTNING_API_KEY"] = apikey
        s = Studio(name=name, teamspace=teamspace, user=user)
        s.start()
        log_content = log_stream.getvalue().strip()
        print(log_content)
        os.environ.pop("LIGHTNING_USER_ID", None)
        os.environ.pop("LIGHTNING_API_KEY", None)
        return {"status": "success", "message": log_content}
    except Exception as e:
        return {"status": "error", "message": str(e)}

@app.post("/stop-studio/")
def stop_studio(userid: str, apikey: str, name: str, teamspace: str, user: str):
    try:
        os.environ["LIGHTNING_USER_ID"] = userid
        os.environ["LIGHTNING_API_KEY"] = apikey
        s = Studio(name=name, teamspace=teamspace, user=user)
        s.stop()
        log_content = log_stream.getvalue().strip()
        print(log_content)
        os.environ.pop("LIGHTNING_USER_ID", None)
        os.environ.pop("LIGHTNING_API_KEY", None)
        return {"status": "success", "message": log_content}
    except Exception as e:
        return {"status": "error", "message": str(e)}