main.py

from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from utils import is_valid_url, bytes_to_human_readable, encode_episodeid
import os
import json
import urllib.parse
from threading import Thread
from LoadBalancer import LoadBalancer
import logging

app = FastAPI()

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

logging.basicConfig(level=logging.INFO)

# Constants and Configuration
CACHE_DIR = os.getenv("CACHE_DIR")
INDEX_FILE = os.getenv("INDEX_FILE")
TOKEN = os.getenv("TOKEN")
REPO = os.getenv("REPO")

load_balancer = LoadBalancer(cache_dir=CACHE_DIR, index_file=INDEX_FILE, token=TOKEN, repo=REPO)

# Start polling in a separate thread
polling_thread = Thread(target=load_balancer.start_polling)
polling_thread.start()

# API Endpoints
@app.get("/api/film/{title}")
async def get_movie_api(title: str):
    """Endpoint to get the movie by title."""
    if not title:
        raise HTTPException(status_code=400, detail="Title parameter is required")
    
    # Check if the film is already cached
    if title in load_balancer.FILM_STORE:
        url = load_balancer.FILM_STORE[title]
        return JSONResponse(content={"url": url})

    movie_path = await load_balancer.find_movie_path(title)
    
    if not movie_path:
        raise HTTPException(status_code=404, detail="Movie not found")
    
    # Start the download in an instance
    response = await load_balancer.download_film_to_best_instance(title=title)
    if response:
        return JSONResponse(content=response)

@app.get("/api/tv/{title}/{season}/{episode}")
async def get_tv_show_api(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 load_balancer.TV_STORE and season in load_balancer.TV_STORE[title]:
        for ep in load_balancer.TV_STORE[title][season]:
            if episode in ep:
                url = load_balancer.TV_STORE[title][season][ep]
                return JSONResponse(content={"url": url})

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

    episode_path = None
    for directory in load_balancer.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")
    
    # Start the download in an instance
    response = await load_balancer.download_episode_to_best_instance(title=title, season=season, episode=episode)
    if response:
        return JSONResponse(content=response)

@app.get("/api/filmid/{title}")
async def get_film_id_by_title_api(title: str):
    """Endpoint to get the film ID by providing the movie title."""
    if not title:
        raise HTTPException(status_code=400, detail="Title parameter is required")
    film_id = await load_balancer.get_film_id(title)
    return JSONResponse(content={"film_id": film_id})

@app.get("/api/episodeid/{title}/{season}/{episode}")
async def get_episode_id_api(title: str, season: str, episode: str):
    """Endpoint to get the episode ID by providing the TV show 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")
    episode_id = encode_episodeid(title, season, episode)
    return JSONResponse(content={"episode_id": episode_id})

@app.get("/api/cache/size")
async def get_cache_size_api():
    total_size = 0
    for dirpath, dirnames, filenames in os.walk(CACHE_DIR):
        for f in filenames:
            fp = os.path.join(dirpath, f)
            total_size += os.path.getsize(fp)
    readable_size = bytes_to_human_readable(total_size)
    return JSONResponse(content={"cache_size": readable_size})

@app.post("/api/cache/clear")
async def clear_cache_api():
    for dirpath, dirnames, filenames in os.walk(CACHE_DIR):
        for f in filenames:
            fp = os.path.join(dirpath, f)
            os.remove(fp)
    return JSONResponse(content={"status": "Cache cleared"})

@app.get("/api/tv/store")
async def get_tv_store_api():
    """Endpoint to get the TV store JSON."""
    return JSONResponse(content=load_balancer.TV_STORE)


@app.get("/api/film/store")
async def get_film_store_api():
    """Endpoint to get the film store JSON."""
    return JSONResponse(content=load_balancer.FILM_STORE)

@app.get("/api/film/metadata/{title}")
async def get_film_metadata_api(title: str):
    """Endpoint to get the film metadata by title."""
    if not title:
        raise HTTPException(status_code=400, detail="No title provided")
    
    json_cache_path = os.path.join(CACHE_DIR, f"{urllib.parse.quote(title)}.json")
    
    if os.path.exists(json_cache_path):
        with open(json_cache_path, 'r') as f:
            data = json.load(f)
        return JSONResponse(content=data)
    
    raise HTTPException(status_code=404, detail="Metadata not found")

@app.get("/api/tv/metadata/{title}")
async def get_tv_metadata_api(title: str):
    """Endpoint to get the TV show metadata by title."""
    if not title:
        raise HTTPException(status_code=400, detail="No title provided")
    
    json_cache_path = os.path.join(CACHE_DIR, f"{urllib.parse.quote(title)}.json")
    
    if os.path.exists(json_cache_path):
        with open(json_cache_path, 'r') as f:
            data = json.load(f)
        
        # Add the file structure to the metadata
        tv_structure_data = await load_balancer.get_tv_structure(title)
        if tv_structure_data:
            data['file_structure'] = tv_structure_data
        
        return JSONResponse(content=data)
    
    raise HTTPException(status_code=404, detail="Metadata not found")

@app.get("/api/film/all")
async def get_all_films_api():
    return JSONResponse(content=await load_balancer.get_all_films())

@app.get("/api/tv/all")
async def get_all_tvshows_api():
    return JSONResponse(content=await load_balancer.get_all_tv_shows())

@app.get("/api/instances")
async def get_instances():
    return JSONResponse(content=load_balancer.instances)

@app.get("/api/instances/health")
async def get_instances_health():
    return JSONResponse(content=load_balancer.instances_health)

# This API is only for instances
@app.post("/api/register")
async def register_instance(request: Request):
    try:
        data = await request.json()
        if not data or "url" not in data:
            raise HTTPException(status_code=400, detail="No URL provided")
        
        url = data["url"]
        if not is_valid_url(url):
            raise HTTPException(status_code=400, detail="Invalid URL")

        # Register the instance
        load_balancer.register_instance(url)
        logging.info(f"Instance registered: {url}")

        return JSONResponse(content={"message": f"Instance {url} registered successfully"}, status_code=200)

    except Exception as e:
        logging.error(f"Error registering instance: {e}")
        raise HTTPException(status_code=500, detail="Failed to register instance")
    
# Routes
@app.get("/")
async def index():
    return f"Load Balancer is Running {load_balancer.version}"