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

# 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/{title}")
async def get_media_api(request: Request, title: str):
    """Endpoint to get the music file (audio or video) by title with support for range requests."""
    if not title:
        raise HTTPException(status_code=400, detail="Title parameter is required")
    
    if title in instance.MUSIC_STORE:
        cache_path = instance.MUSIC_STORE[title]
        if os.path.exists(cache_path):
            return await serve_media(cache_path, request)
    
    media_path = instance.find_music_path(title)
    
    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}"
    media_id = instance.get_music_id(title)
    
    if media_id not in instance.download_threads or not instance.download_threads[media_id].is_alive():
        thread = Thread(target=instance.download_music, args=(file_url, TOKEN, cache_path, media_id, title))
        instance.download_threads[media_id] = thread
        thread.start()
    
    return JSONResponse({"status": "Download started", "music_id": media_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})