app.py

from fastapi.responses import JSONResponse, StreamingResponse, Response
from fastapi.middleware.cors import CORSMiddleware
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()

origins = ["*"]

app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

async def serve_video(file_path: str, request: Request):
    """Serve video file with support for range requests and proper CORS headers."""
    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

    # CORS headers to attach to all responses
    cors_headers = {
        'Access-Control-Allow-Origin': '*',
        'Access-Control-Expose-Headers': 'Content-Length, Content-Range'
    }

    if range_header:
        # Parse range header (e.g., "bytes=START-END")
        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 values
        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")

        # Calculate content length and prepare headers
        content_length = (end - start + 1) if end is not None else file_size - start
        headers = {
            'Content-Range': f'bytes {start}-{end}/{file_size}',
            'Accept-Ranges': 'bytes',
            'Content-Length': str(content_length),
            'Content-Type': mime_type
        }
        headers.update(cors_headers)

        async with aiofiles.open(file_path, 'rb') as f:
            await f.seek(start)
            chunk_size = 8 * 1024 * 1024  # 8 MB
            data = bytearray()
            current = start
            while current <= end:
                remaining = end - current + 1
                read_size = min(chunk_size, remaining)
                chunk = await f.read(read_size)
                if not chunk:
                    break
                data.extend(chunk)
                current += len(chunk)

        return Response(content=bytes(data), status_code=206, headers=headers)
    else:
        # Fallback: Stream the entire file using StreamingResponse with proper headers
        def iterfile():
            with open(file_path, 'rb') as f:
                while True:
                    chunk = f.read(8192)
                    if not chunk:
                        break
                    yield chunk

        headers = {
            'Content-Length': str(file_size),
            'Accept-Ranges': 'bytes',
            'Content-Type': mime_type,
        }
        headers.update(cors_headers)

        return StreamingResponse(iterfile(), media_type=mime_type, headers=headers)

@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 film 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})