"""Serve the original music studio and proxy StepFun using a Space Secret."""
import asyncio
import os
import time
from collections import deque
from pathlib import Path

import httpx
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse, RedirectResponse
from fastapi.staticfiles import StaticFiles

app = FastAPI(docs_url=None, redoc_url=None, openapi_url=None)
UPSTREAM = 'https://api.stepfun.com'
MAX_BODY = 30 * 1024 * 1024
submissions = deque()
submit_lock = asyncio.Lock()


def error(status: int, kind: str, message: str) -> JSONResponse:
    return JSONResponse({'error': {'type': kind, 'message': message}}, status_code=status)


@app.get('/')
async def studio():
    return RedirectResponse('/studio/app.html')


@app.get('/healthz')
async def health() -> dict:
    return {'status': 'ok', 'configured': bool(os.environ.get('STEPFUN_API_KEY'))}


@app.post('/v1/audio/music/{action}')
async def music(action: str, request: Request):
    """Forward only music submission and polling; never return the API key."""
    if action not in ('submit', 'query'):
        return error(404, 'not_found', 'Unknown endpoint')
    key = os.environ.get('STEPFUN_API_KEY', '').strip()
    if not key:
        return error(503, 'service_unavailable', 'Music service is not configured yet')
    raw = bytearray()
    async for chunk in request.stream():
        raw.extend(chunk)
        if len(raw) > MAX_BODY:
            return error(413, 'request_params_invalid', 'Payload too large')
    import json
    try:
        body = json.loads(raw)
        if not isinstance(body, dict):
            raise ValueError()
    except (ValueError, UnicodeDecodeError):
        return error(400, 'request_params_invalid', 'Expected a JSON object')
    if action == 'submit':
        if body.get('task') not in ('text_to_music', 'vocal_to_music', 'music_cover'):
            return error(400, 'request_params_invalid', 'Unsupported task')
        if not isinstance(body.get('caption'), str) or not body['caption'].strip():
            return error(400, 'request_params_invalid', 'Caption is required')
        body['model_id'] = 'step-music'
        async with submit_lock:
            now = time.monotonic()
            while submissions and submissions[0] < now - 60:
                submissions.popleft()
            if len(submissions) >= int(os.environ.get('SUBMITS_PER_MINUTE', '6')):
                return error(429, 'rate_limited', 'Please wait before generating again')
            submissions.append(now)
    elif not isinstance(body.get('task_id'), str) or not body['task_id']:
        return error(400, 'request_params_invalid', 'Task ID is required')
    try:
        async with httpx.AsyncClient(timeout=httpx.Timeout(230, connect=15)) as client:
            res = await client.post(f'{UPSTREAM}/v1/audio/music/{action}', json=body,
                headers={'Authorization': f'Bearer {key}'})
        payload = res.json()
        # Never echo upstream error messages or arbitrary response headers.
        if res.is_error:
            kind = payload.get('error', {}).get('type', 'service_unavailable')
            if kind not in ('not_allowed', 'request_params_invalid', 'model_invalid', 'rate_limited', 'service_unavailable'):
                kind = 'service_unavailable'
            return error(res.status_code, kind, 'Music API request failed')
        return JSONResponse(payload, headers={'Cache-Control': 'no-store'})
    except (httpx.HTTPError, ValueError):
        return error(502, 'service_unavailable', 'Music API is temporarily unavailable')


app.mount('/studio', StaticFiles(directory=Path(__file__).parent / 'static', html=True), name='studio')