app.py

from fastapi.responses import JSONResponse, FileResponse, Response
from fastapi import FastAPI, HTTPException, Request
import mimetypes
from threading import Thread
from Instance import Instance
from api import LoadBalancerAPI
import os
import aiofiles
from urllib.parse import unquote

# Constants and Configuration
CACHE_DIR = os.getenv("CACHE_DIR")
TOKEN = os.getenv("TOKEN")
REPO = os.getenv("REPO")
ID = os.getenv("ID")
URL = os.getenv("URL")
LOAD_BALANCER_URL = os.getenv("LOAD_BALANCER_URL")

load_balancer_api = LoadBalancerAPI(base_url=LOAD_BALANCER_URL)
instance = Instance(id=ID, url=URL, cache_dir=CACHE_DIR, token=TOKEN, repo=REPO, load_balancer_api=load_balancer_api)

app = FastAPI()

from fastapi import HTTPException, Request
from fastapi.responses import Response, FileResponse
import os
import mimetypes
import aiofiles

async def serve_media(file_path: str, request: Request):
    """Serve any media file with support for range requests."""
    if not os.path.isfile(file_path):
        raise HTTPException(status_code=404, detail="Media file not found")
    
    file_size = os.path.getsize(file_path)
    range_header = request.headers.get('range', None)

    # Determine the MIME type based on the file extension
    mime_type, _ = mimetypes.guess_type(file_path)
    if mime_type is None:
        mime_type = 'application/octet-stream'  # Fallback MIME type

    # Handle range requests
    if range_header:
        range_specifier = range_header.replace('bytes=', '').strip()
        start, end = (None, None)

        if '-' in range_specifier:
            start_str, end_str = range_specifier.split('-')
            start = int(start_str)
            end = int(end_str) if end_str else file_size - 1

        # Validate the range
        if start is None or start >= file_size or (end is not None and end >= file_size) or (end is not None and start > end):
            raise HTTPException(status_code=416, detail="Requested range not satisfiable")

        headers = {
            'Content-Range': f'bytes {start}-{end or file_size - 1}/{file_size}',
            'Accept-Ranges': 'bytes',
            'Content-Length': str((end - start + 1) if end is not None else file_size - start),
            'Content-Type': mime_type
        }

        async with aiofiles.open(file_path, 'rb') as f:
            await f.seek(start)
            chunk_size = 8192  # Read in 8KB chunks
            data = bytearray()

            while start <= (end or file_size - 1):
                remaining = (end or file_size - 1) - start + 1
                read_size = min(chunk_size, remaining)
                chunk = await f.read(read_size)

                if not chunk:  # End of file
                    break
                
                data.extend(chunk)
                start += read_size

            return Response(content=bytes(data), status_code=206, headers=headers)

    # Fallback for serving the entire file if no range requested
    return FileResponse(file_path, media_type=mime_type)

@app.get("/")
async def index():
    return instance.version

@app.get("/api/get/report")
async def get_report():
    report = instance.compile_report()
    return JSONResponse(report)

@app.get('/api/get/music/store')
async def get_media_store_api():
    """Endpoint to get the music store JSON."""
    return JSONResponse(instance.MUSIC_STORE)

@app.get("/api/get/music/{filename}")
async def get_media_api(request: Request, filename: str):
    """Endpoint to get the music file (audio or video) by filename with support for range requests."""
    filename = unquote(filename)
    if not filename:
        raise HTTPException(status_code=400, detail="filename parameter is required")
    
    if filename in instance.MUSIC_STORE:
        cache_path = instance.MUSIC_STORE[filename]
        if os.path.exists(cache_path):
            return await serve_media(cache_path, request)
    
    media_path = instance.find_music_path(filename)
    
    if not media_path:
        raise HTTPException(status_code=404, detail="Media not found")
    
    cache_path = os.path.join(CACHE_DIR, media_path)
    file_url = f"https://huggingface.co/{REPO}/resolve/main/{media_path}"
    music_id = instance.get_music_id(filename)
    
    if music_id not in instance.download_threads or not instance.download_threads[music_id].is_alive():
        thread = Thread(target=instance.download_music, args=(file_url, TOKEN, cache_path, music_id, filename))
        instance.download_threads[music_id] = thread
        thread.start()
    
    return JSONResponse({"status": "Download started", "music_id": music_id})

@app.get("/api/get/progress/{id}")
async def get_progress_api(id: str):
    """Endpoint to get the download progress of a media file."""
    progress = instance.get_download_progress(id)
    return JSONResponse({"id": id, "progress": progress})