ohm_streaming/main.py

from fastapi import FastAPI, UploadFile, File, BackgroundTasks, HTTPException
from fastapi.responses import StreamingResponse, FileResponse, JSONResponse, Response
from fastapi.responses import HTMLResponse
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from fastapi import Request
import uvicorn
from pathlib import Path
from typing import List
import shutil
import os
import re

from app.models import DownloadRequest, DownloadTask, DownloadStatus
from app.download_manager import DownloadManager
from app.downloaders import AnimeSamaDownloader
from app import providers

app = FastAPI(title="Ohm Stream Downloader")

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

# Initialize download manager
download_manager = DownloadManager(download_dir="downloads", max_parallel=3)

# Mount static files and templates
app.mount("/static", StaticFiles(directory="static"), name="static")
app.mount("/downloads", StaticFiles(directory="downloads"), name="downloads")
templates = Jinja2Templates(directory="templates")


@app.get("/")
async def root():
    return {
        "message": "Ohm Stream Downloader API",
        "status": "running",
        "version": "2.0",
        "endpoints": {
            "POST /api/download": "Start a new download",
            "GET /api/downloads": "List all downloads",
            "GET /api/download/{task_id}": "Get download status",
            "POST /api/download/{task_id}/pause": "Pause a download",
            "POST /api/download/{task_id}/resume": "Resume a download",
            "DELETE /api/download/{task_id}": "Cancel a download",
            "GET /api/providers": "List all supported providers",
            "GET /web": "Web interface"
        }
    }


@app.get("/api/providers")
async def list_providers():
    """List all supported anime and file hosting providers"""
    return {
        "anime_providers": providers.get_anime_providers(),
        "file_hosts": providers.get_file_hosts()
    }


@app.get("/health")
async def health():
    return {"status": "healthy"}


# Web Interface
@app.get("/web")
async def web_interface(request: Request):
    return templates.TemplateResponse("index.html", {"request": request})


# API Endpoints
@app.post("/api/download")
async def create_download(request: DownloadRequest, background_tasks: BackgroundTasks):
    """Create a new download task"""
    task = download_manager.create_task(request)
    background_tasks.add_task(download_manager.start_download, task.id)
    return {"task_id": task.id, "task": task}


@app.get("/api/downloads")
async def list_downloads():
    """List all download tasks"""
    return {"downloads": download_manager.get_all_tasks()}


@app.get("/api/download/{task_id}")
async def get_download_status(task_id: str):
    """Get status of a specific download"""
    task = download_manager.get_task(task_id)
    if not task:
        raise HTTPException(status_code=404, detail="Task not found")
    return task


@app.post("/api/download/{task_id}/pause")
async def pause_download(task_id: str):
    """Pause a download"""
    task = download_manager.get_task(task_id)
    if not task:
        raise HTTPException(status_code=404, detail="Task not found")
    await download_manager.pause_download(task_id)
    return {"status": "paused"}


@app.post("/api/download/{task_id}/resume")
async def resume_download(task_id: str, background_tasks: BackgroundTasks):
    """Resume a paused download"""
    task = download_manager.get_task(task_id)
    if not task:
        raise HTTPException(status_code=404, detail="Task not found")

    if task.status == DownloadStatus.PAUSED:
        background_tasks.add_task(download_manager.start_download, task_id)
        return {"status": "resumed"}

    return {"status": "already running or completed"}


@app.delete("/api/download/{task_id}")
async def cancel_download(task_id: str):
    """Cancel a download"""
    task = download_manager.get_task(task_id)
    if not task:
        raise HTTPException(status_code=404, detail="Task not found")
    await download_manager.cancel_download(task_id)
    return {"status": "cancelled"}


@app.get("/api/download/{task_id}/file")
async def download_file(task_id: str):
    """Download the completed file"""
    task = download_manager.get_task(task_id)
    if not task:
        raise HTTPException(status_code=404, detail="Task not found")

    if task.status != DownloadStatus.COMPLETED:
        raise HTTPException(status_code=400, detail="Download not completed")

    if not task.file_path or not os.path.exists(task.file_path):
        raise HTTPException(status_code=404, detail="File not found")

    return FileResponse(
        task.file_path,
        filename=task.filename,
        media_type='application/octet-stream'
    )


# Unified Anime Search endpoints
@app.get("/api/anime/search")
async def search_anime_unified(q: str, lang: str = "vostfr"):
    """Search across all anime providers"""
    import time
    import asyncio
    from app.providers import get_anime_providers
    from app.downloaders import AnimeSamaDownloader, AnimeUltimeDownloader, NekoSamaDownloader, VostfreeDownloader

    print(f"\n[SEARCH] Starting search for '{q}' in {lang}")
    start_time = time.time()

    results = {}

    # Create downloader instances
    downloaders = {
        "anime-sama": AnimeSamaDownloader(),
        "anime-ultime": AnimeUltimeDownloader(),
        "neko-sama": NekoSamaDownloader(),
        "vostfree": VostfreeDownloader()
    }

    # Search across all providers in parallel with timeout
    search_tasks = []
    provider_ids = []

    for provider_id, provider in get_anime_providers().items():
        if provider_id in downloaders:
            downloader = downloaders[provider_id]
            print(f"[SEARCH] Queueing search on {provider_id}...")
            search_tasks.append(downloader.search_anime(q, lang))
            provider_ids.append(provider_id)

    # Wait for all searches to complete with a timeout per provider
    print(f"[SEARCH] Waiting for {len(search_tasks)} searches...")
    search_results = await asyncio.gather(*search_tasks, return_exceptions=True)

    # Combine results
    for provider_id, result in zip(provider_ids, search_results):
        if isinstance(result, Exception):
            print(f"[SEARCH] {provider_id} error: {str(result)}")
        elif result:
            print(f"[SEARCH] {provider_id} found {len(result)} results")
            results[provider_id] = result
        else:
            print(f"[SEARCH] {provider_id} no results")

    elapsed = time.time() - start_time
    print(f"[SEARCH] Completed in {elapsed:.2f}s - Total results: {sum(len(r) for r in results.values())}\n")

    return {
        "query": q,
        "lang": lang,
        "results": results
    }


@app.get("/api/anime/episodes")
async def get_anime_episodes(url: str, lang: str = "vostfr"):
    """Get list of episodes for an anime"""
    from app.downloaders import get_downloader

    downloader = get_downloader(url)
    episodes = await downloader.get_episodes(url, lang)

    return {
        "url": url,
        "lang": lang,
        "episodes": episodes
    }


@app.get("/api/anime/providers")
async def get_anime_providers_list():
    """Get list of anime providers with info"""
    from app.providers import get_anime_providers
    return {"providers": get_anime_providers()}


# Anime-Sama specific endpoints (legacy)
@app.get("/api/anime-sama/search")
async def search_anime_sama(q: str, lang: str = "vostfr"):
    """Search for anime on anime-sama"""
    downloader = AnimeSamaDownloader()
    results = await downloader.search_anime(q, lang)
    return {"query": q, "lang": lang, "results": results}


@app.post("/api/anime/download")
async def download_anime_episode(
    url: str,
    background_tasks: BackgroundTasks,
    episode: str | None = None
):
    """Download an anime episode"""
    # Construct episode URL if not provided
    if episode and 'episode-' not in url:
        url = f"{url.rstrip('/')}/episode-{episode}"

    request = DownloadRequest(url=url)
    task = download_manager.create_task(request)
    background_tasks.add_task(download_manager.start_download, task.id)
    return {"task_id": task.id, "task": task}


@app.post("/api/anime/download-season")
async def download_anime_season(
    url: str,
    background_tasks: BackgroundTasks,
    lang: str = "vostfr"
):
    """Download all episodes of an anime season"""
    from app.downloaders import get_downloader

    downloader = get_downloader(url)
    episodes = await downloader.get_episodes(url, lang)

    if not episodes:
        raise HTTPException(status_code=404, detail="No episodes found")

    # Create download tasks for all episodes
    task_ids = []
    for episode in episodes:
        request = DownloadRequest(url=episode['url'])
        task = download_manager.create_task(request)
        task_ids.append(task.id)
        background_tasks.add_task(download_manager.start_download, task.id)

    return {
        "message": f"Started downloading {len(task_ids)} episodes",
        "task_ids": task_ids,
        "total_episodes": len(episodes)
    }


# Video Streaming endpoints
@app.get("/video/{task_id}")
async def stream_video(task_id: str, request: Request):
    """Stream a video file with Range support for seeking"""
    task = download_manager.get_task(task_id)
    if not task:
        raise HTTPException(status_code=404, detail="Task not found")

    if task.status != DownloadStatus.COMPLETED:
        raise HTTPException(status_code=400, detail="Download not completed")

    if not task.file_path or not os.path.exists(task.file_path):
        raise HTTPException(status_code=404, detail="File not found")

    file_path = Path(task.file_path)
    file_size = file_path.stat().st_size

    # Parse Range header
    range_header = request.headers.get("range")
    headers = {
        "Accept-Ranges": "bytes",
        "Content-Type": "video/mp4",
    }

    if range_header:
        # Parse Range header (format: bytes=start-end)
        try:
            range_match = re.match(r"bytes=(\d+)-(\d*)", range_header)
            start = int(range_match.group(1))
            end = int(range_match.group(2)) if range_match.group(2) else file_size - 1

            # Validate range
            if start >= file_size or end >= file_size or start > end:
                headers["Content-Range"] = f"bytes */{file_size}"
                return Response(
                    status_code=416,
                    headers=headers,
                    content="Requested Range Not Satisfiable"
                )

            # Read the requested range
            content_length = end - start + 1
            headers["Content-Range"] = f"bytes {start}-{end}/{file_size}"
            headers["Content-Length"] = str(content_length)

            async def video_range_reader():
                with open(file_path, 'rb') as f:
                    f.seek(start)
                    remaining = content_length
                    while remaining > 0:
                        chunk_size = min(1024 * 1024, remaining)  # 1MB chunks
                        data = f.read(chunk_size)
                        if not data:
                            break
                        remaining -= len(data)
                        yield data

            return Response(
                content=video_range_reader(),
                status_code=206,
                headers=headers
            )
        except Exception as e:
            raise HTTPException(status_code=400, detail=f"Invalid Range header: {e}")
    else:
        # No Range header - stream entire file
        async def video_reader():
            with open(file_path, 'rb') as f:
                while True:
                    data = f.read(1024 * 1024)  # 1MB chunks
                    if not data:
                        break
                    yield data

        headers["Content-Length"] = str(file_size)
        return Response(
            content=video_reader(),
            headers=headers
        )


# Direct video streaming endpoint (by filename)
@app.get("/stream/{filename}")
async def stream_video_by_filename(filename: str, request: Request):
    """Stream a video file by filename with Range support for seeking"""
    # Sanitize filename to prevent directory traversal
    filename = os.path.basename(filename)
    file_path = Path("downloads") / filename

    if not file_path.exists():
        raise HTTPException(status_code=404, detail="File not found")

    file_size = file_path.stat().st_size

    # Parse Range header
    range_header = request.headers.get("range")

    if range_header:
        # Parse Range header (format: bytes=start-end)
        try:
            range_match = re.match(r"bytes=(\d+)-(\d*)", range_header)
            start = int(range_match.group(1))
            end = int(range_match.group(2)) if range_match.group(2) else file_size - 1

            # Validate range
            if start >= file_size or end >= file_size or start > end:
                return Response(
                    status_code=416,
                    headers={
                        "Content-Range": f"bytes */{file_size}",
                        "Accept-Ranges": "bytes"
                    },
                    content="Requested Range Not Satisfiable"
                )

            # Read the requested range
            content_length = end - start + 1

            def video_range_reader():
                with open(file_path, 'rb') as f:
                    f.seek(start)
                    remaining = content_length
                    while remaining > 0:
                        chunk_size = min(1024 * 1024, remaining)  # 1MB chunks
                        data = f.read(chunk_size)
                        if not data:
                            break
                        remaining -= len(data)
                        yield data

            return StreamingResponse(
                video_range_reader(),
                status_code=206,
                headers={
                    "Content-Range": f"bytes {start}-{end}/{file_size}",
                    "Content-Length": str(content_length),
                    "Accept-Ranges": "bytes",
                    "Content-Type": "video/mp4",
                }
            )
        except Exception as e:
            raise HTTPException(status_code=400, detail=f"Invalid Range header: {e}")
    else:
        # No Range header - stream entire file
        def video_reader():
            with open(file_path, 'rb') as f:
                while True:
                    data = f.read(1024 * 1024)  # 1MB chunks
                    if not data:
                        break
                    yield data

        return StreamingResponse(
            video_reader(),
            headers={
                "Content-Length": str(file_size),
                "Accept-Ranges": "bytes",
                "Content-Type": "video/mp4",
            }
        )


# Video Player page (by task_id)
@app.get("/player/{task_id}")
async def video_player(request: Request, task_id: str):
    """Video player page for watching downloaded anime"""
    task = download_manager.get_task(task_id)
    if not task:
        raise HTTPException(status_code=404, detail="Task not found")

    if task.status != DownloadStatus.COMPLETED:
        raise HTTPException(status_code=400, detail="Download not completed")

    if not task.file_path or not os.path.exists(task.file_path):
        raise HTTPException(status_code=404, detail="File not found")

    # Get video info
    file_path = Path(task.file_path)
    file_size = file_path.stat().st_size

    # Calculate video duration (rough estimation based on file size)
    # Assuming ~1MB per minute for 720p, ~2MB per minute for 1080p
    estimated_duration_seconds = int(file_size / (1.5 * 1024 * 1024))

    return templates.TemplateResponse("player.html", {
        "request": request,
        "task_id": task_id,
        "filename": task.filename,
        "file_size": file_size,
        "estimated_duration": estimated_duration_seconds
    })


# Video Player page (by filename)
@app.get("/watch/{filename}")
async def video_player_by_filename(request: Request, filename: str):
    """Video player page for watching downloaded anime by filename"""
    # Sanitize filename
    filename = os.path.basename(filename)
    file_path = Path("downloads") / filename

    if not file_path.exists():
        raise HTTPException(status_code=404, detail="File not found")

    file_size = file_path.stat().st_size
    estimated_duration_seconds = int(file_size / (1.5 * 1024 * 1024))

    return templates.TemplateResponse("player.html", {
        "request": request,
        "task_id": filename,  # Use filename instead of task_id
        "filename": filename,
        "file_size": file_size,
        "estimated_duration": estimated_duration_seconds
    })


if __name__ == "__main__":
    uvicorn.run(
        "main:app",
        host="0.0.0.0",
        port=3000,
        reload=True
    )