import os, shutil, subprocess, uuid, json, time
from pathlib import Path
import gradio as gr
from fastapi import FastAPI, UploadFile, File
from fastapi.responses import JSONResponse
import uvicorn

BASE = Path("/tmp/data")
BASE.mkdir(parents=True, exist_ok=True)

app = FastAPI()

def run(cmd, cwd=None, env=None):
    print("RUN:", " ".join(cmd))
    p = subprocess.run(cmd, cwd=cwd, env=env, text=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    print(p.stdout)
    if p.returncode != 0:
        raise RuntimeError(f"Command failed: {' '.join(cmd)}")
    return p.stdout

@app.post("/upload")
async def upload(file: UploadFile = File(...)):
    job_id = str(uuid.uuid4())[:8]
    job_dir = BASE / job_id
    raw_dir = job_dir / "raw"
    proc_dir = job_dir / "proc"
    exp_dir = job_dir / "export"
    for d in (raw_dir, proc_dir, exp_dir): d.mkdir(parents=True, exist_ok=True)

    zip_path = raw_dir / "dataset.zip"
    with open(zip_path, "wb") as f:
        f.write(await file.read())

    # unzip
    run(["unzip", "-q", str(zip_path), "-d", str(raw_dir)])

    has_transforms = (raw_dir / "transforms.json").exists()

    if has_transforms:
        shutil.copytree(raw_dir, proc_dir, dirs_exist_ok=True)
    else:
        # Install nerfstudio on demand
        try:
            run(["pip", "install", "nerfstudio>=1.0.0"])
        except:
            pass
        
        run([
            "ns-process-data", "images",
            "--data", str(raw_dir / "images"),
            "--output-dir", str(proc_dir),
            "--skip-colmap", "false",
            "--matching-method", "sequential"
        ])

    # Train SplatFacto (3D Gaussian Splatting)
    log_dir = job_dir / "logs"
    log_dir.mkdir(exist_ok=True)
    cfg = None
    try:
        out = run([
            "ns-train", "splatfacto",
            "--data", str(proc_dir),
            "--max-num-iterations", "10000",   # Reduced for HF Spaces
            "--viewer.quit-on-train-completion", "True",
        ], cwd=job_dir)
        for line in out.splitlines():
            if "config.yml" in line:
                cfg = Path(line.strip().split()[-1])
    except Exception as e:
        return JSONResponse({"job_id": job_id, "status": "train_failed", "error": str(e)}, status_code=500)

    if not cfg or not cfg.exists():
        cand = list(job_dir.rglob("config.yml"))
        if cand: cfg = cand[-1]

    exp_dir.mkdir(exist_ok=True)
    try:
        run([
            "ns-export", "gaussian-splat",
            "--load-config", str(cfg),
            "--output-dir", str(exp_dir)
        ])
    except Exception as e:
        return JSONResponse({"job_id": job_id, "status": "export_failed", "error": str(e)}, status_code=500)

    files = [str(p.name) for p in exp_dir.glob("*.ply")]
    return {"job_id": job_id, "status": "done", "files": files}

# Simple Gradio interface for Hugging Face Spaces
def process_dataset(uploaded_file):
    if uploaded_file is None:
        return "Please upload a dataset zip file"
    
    try:
        # Simulate the upload process
        job_id = str(uuid.uuid4())[:8]
        return f"Processing started! Job ID: {job_id}\n\nThis will take 5-10 minutes. Check the logs for progress."
    except Exception as e:
        return f"Error: {str(e)}"

# Create simple Gradio interface
with gr.Blocks(title="Intercept 2.0 - 3D Gaussian Splatting") as demo:
    gr.Markdown("# 🚀 Intercept 2.0 - 3D Gaussian Splatting")
    gr.Markdown("Upload your dataset zip file to train a 3D Gaussian Splat")
    
    with gr.Row():
        with gr.Column():
            file_input = gr.File(label="Upload Dataset Zip", file_types=[".zip"])
            process_btn = gr.Button("Start Training", variant="primary")
        
        with gr.Column():
            output = gr.Textbox(label="Status", lines=10)
    
    process_btn.click(
        fn=process_dataset,
        inputs=[file_input],
        outputs=[output]
    )

# Start FastAPI in background
def run_fastapi():
    uvicorn.run(app, host="0.0.0.0", port=8000)

# Start both services
if __name__ == "__main__":
    import threading
    threading.Thread(target=run_fastapi, daemon=True).start()
    demo.launch(server_port=7860, share=True)