app.py

import os
import asyncio
import tempfile
import shutil
import time
from datetime import datetime, timedelta
from typing import Dict, List, Optional
from fastapi import FastAPI, HTTPException, BackgroundTasks
from fastapi.responses import FileResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from torrentp import TorrentDownloader

app = FastAPI(title="Torrent Downloader Service", version="1.0.0")

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

# Data models
class DownloadRequest(BaseModel):
    magnet_link: str

class DownloadResponse(BaseModel):
    download_id: str
    status: str

class StatusResponse(BaseModel):
    status: str
    progress: Optional[int] = None
    files: Optional[List[str]] = None
    message: Optional[str] = None

class FileInfo(BaseModel):
    filename: str
    size: int
    path: str

class FilesResponse(BaseModel):
    files: List[FileInfo]

# Global storage
download_status: Dict[str, Dict] = {}
download_files: Dict[str, Dict] = {}

def cleanup_old_files():
    """Clean up files older than 1 hour"""
    current_time = datetime.now()
    to_remove = []
    
    for download_id, info in download_files.items():
        if current_time - info['created'] > timedelta(hours=1):
            try:
                if os.path.exists(info['path']):
                    if os.path.isdir(info['path']):
                        shutil.rmtree(info['path'])
                    else:
                        os.remove(info['path'])
                to_remove.append(download_id)
            except Exception as e:
                print(f"Error cleaning up {info['path']}: {e}")
    
    for download_id in to_remove:
        download_files.pop(download_id, None)
        download_status.pop(download_id, None)

async def download_torrent_async(magnet_link: str, download_id: str, temp_dir: str):
    """Async function to download torrent"""
    try:
        download_status[download_id] = {"status": "downloading", "progress": 0}
        
        # Initialize TorrentDownloader
        torrent_downloader = TorrentDownloader(magnet_link, temp_dir, stop_after_download=True)
        await torrent_downloader.start_download()
        
        # Find downloaded files
        downloaded_files = []
        for root, dirs, files in os.walk(temp_dir):
            for file in files:
                file_path = os.path.join(root, file)
                downloaded_files.append(file_path)
        
        if downloaded_files:
            download_status[download_id] = {"status": "completed", "files": downloaded_files}
            # Store file info with creation time for cleanup
            download_files[download_id] = {
                "path": temp_dir,
                "files": downloaded_files,
                "created": datetime.now()
            }
        else:
            download_status[download_id] = {"status": "error", "message": "No files downloaded"}
            
    except Exception as e:
        download_status[download_id] = {"status": "error", "message": str(e)}

@app.post("/download", response_model=DownloadResponse)
async def start_download(request: DownloadRequest, background_tasks: BackgroundTasks):
    """Start a torrent download"""
    cleanup_old_files()  # Clean up old files before starting new download
    
    magnet_link = request.magnet_link
    download_id = str(int(time.time() * 1000))  # Use timestamp as ID
    
    # Create temporary directory for this download
    temp_dir = tempfile.mkdtemp(prefix=f"torrent_{download_id}_")
    
    # Start download in background
    background_tasks.add_task(download_torrent_async, magnet_link, download_id, temp_dir)
    
    download_status[download_id] = {"status": "starting"}
    
    return DownloadResponse(download_id=download_id, status="started")

@app.get("/status/{download_id}", response_model=StatusResponse)
async def get_status(download_id: str):
    """Get download status"""
    if download_id not in download_status:
        raise HTTPException(status_code=404, detail="Download ID not found")
    
    status_data = download_status[download_id]
    return StatusResponse(**status_data)

@app.get("/files/{download_id}", response_model=FilesResponse)
async def list_files(download_id: str):
    """List downloaded files"""
    if download_id not in download_files:
        raise HTTPException(status_code=404, detail="Download ID not found or no files available")
    
    files_info = []
    for file_path in download_files[download_id]["files"]:
        if os.path.exists(file_path):
            files_info.append(FileInfo(
                filename=os.path.basename(file_path),
                size=os.path.getsize(file_path),
                path=file_path
            ))
    
    return FilesResponse(files=files_info)

@app.get("/download-file/{download_id}/{filename}")
async def download_file(download_id: str, filename: str):
    """Download a specific file"""
    if download_id not in download_files:
        raise HTTPException(status_code=404, detail="Download ID not found")
    
    # Find the file
    target_file = None
    for file_path in download_files[download_id]["files"]:
        if os.path.basename(file_path) == filename:
            target_file = file_path
            break
    
    if not target_file or not os.path.exists(target_file):
        raise HTTPException(status_code=404, detail="File not found")
    
    return FileResponse(
        path=target_file,
        filename=filename,
        media_type='application/octet-stream'
    )

@app.delete("/cleanup/{download_id}")
async def cleanup_download(download_id: str):
    """Manually cleanup a download"""
    if download_id in download_files:
        try:
            temp_dir = download_files[download_id]["path"]
            if os.path.exists(temp_dir):
                shutil.rmtree(temp_dir)
            download_files.pop(download_id, None)
            download_status.pop(download_id, None)
            return {"message": "Cleanup successful"}
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"Cleanup failed: {str(e)}")
    else:
        raise HTTPException(status_code=404, detail="Download ID not found")

@app.get("/")
async def root():
    """Root endpoint with service info"""
    return {
        "service": "Torrent Downloader Service",
        "version": "1.0.0",
        "endpoints": {
            "POST /download": "Start a torrent download",
            "GET /status/{download_id}": "Check download status",
            "GET /files/{download_id}": "List downloaded files",
            "GET /download-file/{download_id}/{filename}": "Download a specific file",
            "DELETE /cleanup/{download_id}": "Clean up a download",
            "GET /docs": "API documentation"
        }
    }

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)