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 re
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()

async def serve_video(file_path: str, request: Request):
    """Serve video file with support for range requests."""
    if not os.path.isfile(file_path):
        raise HTTPException(status_code=404, detail="Video 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

    if range_header:
        # Simple parsing for range values
        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 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:  # EOF
                    break
                
                data.extend(chunk)
                start += read_size

            return Response(content=bytes(data), status_code=206, headers=headers)  # Convert bytearray to bytes

    # Fallback for serving the whole 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/tv/store')
async def get_tv_store_api():
    """Endpoint to get the TV store JSON."""
    return JSONResponse(instance.TV_STORE)

@app.get('/api/get/film/store')
async def get_film_store_api():
    """Endpoint to get the TV store JSON."""
    return JSONResponse(instance.FILM_STORE)

@app.get("/api/get/film/{title}")
async def get_movie_api(request: Request, title: str):
    """Endpoint to get the movie by title with support for range requests."""
    if not title:
        raise HTTPException(status_code=400, detail="Title parameter is required")
    
    # Check if the film is already cached
    if title in instance.FILM_STORE:
        cache_path = instance.FILM_STORE[title]
        if os.path.exists(cache_path):
            return await serve_video(cache_path, request)
    
    movie_path = instance.find_movie_path(title)
    
    if not movie_path:
        raise HTTPException(status_code=404, detail="Movie not found")
    
    cache_path = os.path.join(CACHE_DIR, movie_path)
    file_url = f"https://huggingface.co/{REPO}/resolve/main/{movie_path}"
    film_id = instance.get_film_id(title)
    
    # Start the download in a separate thread if not already downloading
    if film_id not in instance.download_threads or not instance.download_threads[film_id].is_alive():
        thread = Thread(target=instance.download_film, args=(file_url, TOKEN, cache_path, film_id, title))
        instance.download_threads[film_id] = thread
        thread.start()
    
    return JSONResponse({"status": "Download started", "film_id": film_id})

@app.get("/api/get/tv/{title}/{season}/{episode}")
async def get_tv_show_api(request: Request, title: str, season: str, episode: str):
    """Endpoint to get the TV show by title, season, and episode."""
    if not title or not season or not episode:
        raise HTTPException(status_code=400, detail="Title, season, and episode parameters are required")

    # Check if the episode is already cached
    if title in instance.TV_STORE and season in instance.TV_STORE[title]:
        for ep in instance.TV_STORE[title][season]:
            if episode in ep:
                cache_path = instance.TV_STORE[title][season][ep]
                if os.path.exists(cache_path):
                    return await serve_video(cache_path, request)

    tv_path = instance.find_tv_path(title)
    
    if not tv_path:
        raise HTTPException(status_code=404, detail="TV show not found")

    episode_path = None
    for directory in instance.file_structure:
        if directory['type'] == 'directory' and directory['path'] == 'tv':
            for sub_directory in directory['contents']:
                if sub_directory['type'] == 'directory' and title.lower() in sub_directory['path'].lower():
                    for season_dir in sub_directory['contents']:
                        if season_dir['type'] == 'directory' and season in season_dir['path']:
                            for episode_file in season_dir['contents']:
                                if episode_file['type'] == 'file' and episode in episode_file['path']:
                                    episode_path = episode_file['path']
                                    break

    if not episode_path:
        raise HTTPException(status_code=404, detail="Episode not found")
    
    cache_path = os.path.join(CACHE_DIR, episode_path)
    file_url = f"https://huggingface.co/{REPO}/resolve/main/{episode_path}"
    episode_id = instance.encode_episodeid(title, season, episode)
    
    # Start the download in a separate thread if not already downloading
    if episode_id not in instance.download_threads or not instance.download_threads[episode_id].is_alive():
        thread = Thread(target=instance.download_episode, args=(file_url, TOKEN, cache_path, episode_id, title))
        instance.download_threads[episode_id] = thread
        thread.start()
    
    return JSONResponse({"status": "Download started", "episode_id": episode_id})

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