import subprocess
import uuid
from typing import Dict, List
from fastapi import APIRouter, HTTPException
from schemas import ServeRequest, ServeStatus

router = APIRouter()

# In-memory store for serving processes
_serving_processes = {}  # type: Dict[str, Dict]


@router.post("/serve", response_model=ServeStatus)
def start_serve(req: ServeRequest):
    serve_id = uuid.uuid4().hex[:8]
    cmd = [
        "mlflow", "models", "serve",
        "-m", req.model_uri,
        "-p", str(req.port),
        "--no-conda",
    ]
    if req.tracking_uri:
        cmd.extend(["--env-manager", "local"])

    try:
        proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    except FileNotFoundError:
        raise HTTPException(status_code=500, detail="mlflow CLI not found")

    _serving_processes[serve_id] = {
        "id": serve_id,
        "model_uri": req.model_uri,
        "port": req.port,
        "pid": proc.pid,
        "process": proc,
    }

    return ServeStatus(
        id=serve_id,
        model_uri=req.model_uri,
        port=req.port,
        pid=proc.pid,
        status="running",
    )


@router.get("/serve", response_model=List[ServeStatus])
def list_serve():
    results = []
    for sid, info in _serving_processes.items():
        proc = info["process"]
        status = "running" if proc.poll() is None else "stopped"
        results.append(ServeStatus(
            id=sid,
            model_uri=info["model_uri"],
            port=info["port"],
            pid=info["pid"],
            status=status,
        ))
    return results


@router.delete("/serve/{serve_id}")
def stop_serve(serve_id: str):
    if serve_id not in _serving_processes:
        raise HTTPException(status_code=404, detail="Serve process not found")

    proc = _serving_processes[serve_id]["process"]
    if proc.poll() is None:
        proc.terminate()

    model_uri = _serving_processes.pop(serve_id)["model_uri"]
    return {"message": f"Stopped serving {model_uri}"}