import os
import sys
import asyncio
import uuid
import threading
from typing import List, Optional

os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"

from fastapi import FastAPI, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from acestep.pipeline_ace_step import ACEStepPipeline
from contextlib import asynccontextmanager
import uvicorn
from concurrent.futures import ThreadPoolExecutor
import logging

# Constants & config
SHUTDOWN_AUTH_TOKEN = "your-secret-token"  # Change to a secure key
AUDIO_OUTPUT_DIR = os.path.join(os.getcwd(), "generated_audio")

# Ensure output directory exists
os.makedirs(AUDIO_OUTPUT_DIR, exist_ok=True)

PRELOAD_CONFIG = {
    "checkpoint_path": "",
    "bf16": True,
    "torch_compile": False,
    "device_id": 0,
}

# Globals
preloaded_pipeline: Optional[ACEStepPipeline] = None
pipeline_lock: Optional[asyncio.Lock] = None
server: Optional[uvicorn.Server] = None
thread_pool: Optional[ThreadPoolExecutor] = ThreadPoolExecutor(max_workers=4)  # Tune max_workers as needed

class ACEStepInput(BaseModel):
    audio_duration: float
    prompt: str
    lyrics: str
    infer_step: int
    guidance_scale: float
    scheduler_type: str
    cfg_type: str
    omega_scale: float
    actual_seeds: List[int]
    guidance_interval: float
    guidance_interval_decay: float
    min_guidance_scale: float
    use_erg_tag: bool
    use_erg_lyric: bool
    use_erg_diffusion: bool
    oss_steps: List[int]
    guidance_scale_text: float = 0.0
    guidance_scale_lyric: float = 0.0

class ACEStepOutput(BaseModel):
    status: str
    output_path: Optional[str]
    message: str

def set_cuda_device(device_id: int) -> None:
    current_cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES")
    device_str = str(device_id)
    if current_cuda_visible_devices != device_str:
        os.environ["CUDA_VISIBLE_DEVICES"] = device_str

def initialize_pipeline(
    checkpoint_path: str,
    bf16: bool,
    torch_compile: bool,
    device_id: int,
) -> ACEStepPipeline:
    set_cuda_device(device_id)
    dtype = "bfloat16" if bf16 else "float32"
    pipeline = ACEStepPipeline(
        checkpoint_dir=checkpoint_path,
        dtype=dtype,
        torch_compile=torch_compile,
    )
    return pipeline

async def run_generation_in_thread(pipeline: ACEStepPipeline, params: tuple, output_path: str):
    loop = asyncio.get_running_loop()
    await loop.run_in_executor(thread_pool, lambda: pipeline(*params, save_path=output_path))

@asynccontextmanager
async def lifespan(app: FastAPI):
    global preloaded_pipeline, pipeline_lock
    cfg = PRELOAD_CONFIG
    preloaded_pipeline = initialize_pipeline(
        cfg["checkpoint_path"],
        cfg["bf16"],
        cfg["torch_compile"],
        cfg["device_id"],
    )
    pipeline_lock = asyncio.Lock()
    print("Preloaded model pipeline at startup.")
    yield
    # Cleanup
    preloaded_pipeline = None
    pipeline_lock = None
    thread_pool.shutdown(wait=False)
    print("Pipeline cleaned up and thread pool shutdown on server shutdown.")

app = FastAPI(title="ACEStep Pipeline API", lifespan=lifespan)

# CORS middleware (allow localhost + LAN IPs or all for development)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # For development: restrict to your Android device IP or LAN domains in production
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Mount static files directory for output
app.mount("/files", StaticFiles(directory=AUDIO_OUTPUT_DIR), name="files")

@app.post("/generate", response_model=ACEStepOutput)
async def generate_audio(input_data: ACEStepInput):
    if preloaded_pipeline is None or pipeline_lock is None:
        raise HTTPException(status_code=503, detail="Model pipeline is not ready")

    try:
        actual_seeds_str = ", ".join(map(str, input_data.actual_seeds))
        oss_steps_str = ", ".join(map(str, input_data.oss_steps))

        params = (
            input_data.audio_duration,
            input_data.prompt,
            input_data.lyrics,
            input_data.infer_step,
            input_data.guidance_scale,
            input_data.scheduler_type,
            input_data.cfg_type,
            input_data.omega_scale,
            actual_seeds_str,
            input_data.guidance_interval,
            input_data.guidance_interval_decay,
            input_data.min_guidance_scale,
            input_data.use_erg_tag,
            input_data.use_erg_lyric,
            input_data.use_erg_diffusion,
            oss_steps_str,
            input_data.guidance_scale_text,
            input_data.guidance_scale_lyric,
        )
        filename = f"output_{uuid.uuid4().hex}.wav"
        output_path = os.path.join(AUDIO_OUTPUT_DIR, filename)

        async with pipeline_lock:
            await run_generation_in_thread(preloaded_pipeline, params, output_path)

        return ACEStepOutput(
            status="success",
            output_path=f"files/{filename}",
            message="Audio generated successfully",
        )
    except Exception as e:
        logging.exception("Exception during generation")
        raise HTTPException(status_code=500, detail=f"Error generating audio: {e}")

@app.get("/health")
async def health_check():
    return {"status": "healthy"}

@app.post("/shutdown")
async def shutdown(request: Request):
    global server
    token = request.headers.get("X-Auth-Token")
    if token != SHUTDOWN_AUTH_TOKEN:
        raise HTTPException(status_code=401, detail="Unauthorized")

    if server is None:
        raise HTTPException(status_code=503, detail="Server instance not found")

    def stop_server():
        print("Shutdown initiated via /shutdown endpoint.")
        server.should_exit = True

    threading.Thread(target=stop_server).start()
    return {"status": "shutting down"}

def run_uvicorn(port: int = 8000):
    global server
    config = uvicorn.Config(
        app,
        host="0.0.0.0",
        port=port,
        log_level="info",
        # workers=1 to avoid concurrency issues if needed
        workers=1,
    )
    server = uvicorn.Server(config)
    print(f"Starting server on 0.0.0.0:{port}")
    server.run()

if __name__ == "__main__":
    port = 8000
    if "--port" in sys.argv:
        try:
            idx = sys.argv.index("--port")
            if idx + 1 < len(sys.argv):
                port = int(sys.argv[idx + 1])
        except Exception:
            print("Invalid port argument passed; falling back to default port 8000")

    run_uvicorn(port)

my uvicorn server cannot be accessed outside of my PC why? if its the script fix and return in full otherwise explain.