app.py

import asyncio
import os
import re
import time
import uuid
import shutil
import aiohttp
import gc  # Memory management ke liye garbage collector
import zipfile  # ZIP file handle karne ke liye in-built library
from fastapi import FastAPI, BackgroundTasks, Request, File, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, StreamingResponse
from pydantic import BaseModel

app = FastAPI()

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

tasks = {}
DATA_DIR = "/data"
os.makedirs(DATA_DIR, exist_ok=True)

class DownloadRequest(BaseModel):
    url: str

async def cleanup_old_files():
    while True:
        try:
            now = time.time()
            for filename in os.listdir(DATA_DIR):
                filepath = os.path.join(DATA_DIR, filename)
                if os.path.isfile(filepath):
                    if now - os.path.getmtime(filepath) > 86400:
                        os.remove(filepath)
                        for tid, tinfo in list(tasks.items()):
                            if tinfo.get("file_path") == filepath:
                                del tasks[tid]
        except Exception as e:
            print(f"Cleanup error: {e}")
        await asyncio.sleep(3600)

@app.on_event("startup")
async def startup_event():
    asyncio.create_task(cleanup_old_files())

async def download_file(task_id: str, url: str):
    file_path = tasks[task_id]["file_path"]
    try:
        connector = aiohttp.TCPConnector(limit=10)
        async with aiohttp.ClientSession(connector=connector, read_bufsize=256 * 1024) as session:
            async with session.get(url, timeout=None) as response:
                response.raise_for_status()
                total_size = int(response.headers.get('Content-Length', 0))
                tasks[task_id]["total_size"] = total_size
                
                cd = response.headers.get('Content-Disposition')
                if cd and 'filename=' in cd:
                    fname = re.findall('filename="([^"]+)"', cd)
                    if not fname:
                        fname = re.findall('filename=([^;]+)', cd)
                    if fname:
                        new_filename = fname[0]
                        new_file_path = os.path.join(DATA_DIR, f"{task_id}_{new_filename}")
                        tasks[task_id]["file_path"] = new_file_path
                        tasks[task_id]["original_filename"] = new_filename
                        file_path = new_file_path

                downloaded = 0
                start_time = time.time()
                last_time = start_time
                last_downloaded = 0
                chunk_counter = 0
                
                with open(file_path, 'wb') as f:
                    async for chunk in response.content.iter_chunked(256 * 1024):
                        if not chunk:
                            break
                        f.write(chunk)
                        downloaded += len(chunk)
                        tasks[task_id]["downloaded"] = downloaded
                        
                        current_time = time.time()
                        if current_time - last_time >= 1.0:
                            speed = (downloaded - last_downloaded) / (current_time - last_time)
                            tasks[task_id]["speed"] = speed
                            last_time = current_time
                            last_downloaded = downloaded
                        
                        f.flush()
                        chunk_counter += 1
                        if chunk_counter % 20 == 0:
                            os.fsync(f.fileno())
                            del chunk
                            gc.collect()
                            
                tasks[task_id]["status"] = "completed"
                tasks[task_id]["speed"] = 0.0
    except Exception as e:
        tasks[task_id]["status"] = "error"
        tasks[task_id]["error"] = str(e)
    finally:
        gc.collect()

@app.post("/start_download")
async def start_download(req: DownloadRequest, background_tasks: BackgroundTasks):
    task_id = str(uuid.uuid4())
    filename = req.url.split("/")[-1]
    if "?" in filename:
        filename = filename.split("?")[0]
    if not filename or "." not in filename:
        filename = "downloaded_file.bin"
        
    tasks[task_id] = {
        "task_id": task_id,
        "url": req.url,
        "status": "downloading",
        "total_size": 0,
        "downloaded": 0,
        "speed": 0.0,
        "file_path": os.path.join(DATA_DIR, f"{task_id}_{filename}"),
        "original_filename": filename,
        "timestamp": time.time()
    }
    background_tasks.add_task(download_file, task_id, req.url)
    return {"task_id": task_id}

@app.get("/status/{task_id}")
async def get_status(task_id: str):
    if task_id not in tasks:
        return {"error": "Task not found"}
    return tasks[task_id]

@app.get("/download/{task_id}")
async def download(task_id: str):
    if task_id not in tasks or tasks[task_id]["status"] != "completed":
        return {"error": "File not ready"}
    return FileResponse(tasks[task_id]["file_path"], filename=tasks[task_id].get("original_filename"))

@app.get("/stream/{task_id}")
async def stream(task_id: str, request: Request):
    if task_id not in tasks or tasks[task_id]["status"] != "completed":
        return {"error": "File not ready"}
    
    file_path = tasks[task_id]["file_path"]
    file_size = os.path.getsize(file_path)
    range_header = request.headers.get("Range")
    
    if range_header:
        byte1, byte2 = 0, None
        match = range_header.replace("bytes=", "").split("-")
        if match[0]:
            byte1 = int(match[0])
        if len(match) > 1 and match[1]:
            byte2 = int(match[1])
        
        length = file_size - byte1
        if byte2 is not None:
            length = byte2 + 1 - byte1
            
        def file_iterator(start, length):
            with open(file_path, "rb") as f:
                f.seek(start)
                chunk_size = 512 * 1024  # Reduced chunk size for memory safety
                while length > 0:
                    read_size = min(chunk_size, length)
                    data = f.read(read_size)
                    if not data:
                        break
                    yield data
                    length -= len(data)
                    del data
                    
        headers = {
            "Content-Range": f"bytes {byte1}-{byte1+length-1}/{file_size}",
            "Accept-Ranges": "bytes",
            "Content-Length": str(length),
        }
        return StreamingResponse(file_iterator(byte1, length), status_code=206, headers=headers)
    else:
        return FileResponse(file_path)

@app.get("/storage")
async def get_storage():
    total, used, free = shutil.disk_usage(DATA_DIR)
    return {"total": total, "used": used, "free": free}

@app.post("/delete_all")
async def delete_all():
    for filename in os.listdir(DATA_DIR):
        filepath = os.path.join(DATA_DIR, filename)
        try:
            if os.path.isfile(filepath):
                os.remove(filepath)
            elif os.path.isdir(filepath):
                shutil.rmtree(filepath)
        except Exception as e:
            print(f"Failed to delete {filepath}: {e}")
    tasks.clear()
    gc.collect()
    return {"status": "success"}

@app.get("/history")
async def get_history():
    return {"history": list(tasks.values())}

@app.post("/upload")
async def upload_file(file: UploadFile = File(...)):
    task_id = str(uuid.uuid4())
    original_filename = file.filename if file.filename else "uploaded_file.bin"
    file_path = os.path.join(DATA_DIR, f"{task_id}_{original_filename}")
    
    # Using chunked reading directly from the file stream to save memory
    with open(file_path, "wb") as buffer:
        while True:
            chunk = await file.read(256 * 1024)
            if not chunk:
                break
            buffer.write(chunk)
            buffer.flush()
            os.fsync(buffer.fileno())
            del chunk
    
    await file.close() # Strictly close file to free memory
    file_size = os.path.getsize(file_path)
    
    tasks[task_id] = {
        "task_id": task_id,
        "url": "local_upload",
        "status": "completed",
        "total_size": file_size,
        "downloaded": file_size,
        "speed": 0.0,
        "file_path": file_path,
        "original_filename": original_filename,
        "timestamp": time.time()
    }
    gc.collect()
    return {"task_id": task_id}

@app.post("/start_gofile_transfer")
async def start_gofile_transfer(req: DownloadRequest, background_tasks: BackgroundTasks):
    task_id = str(uuid.uuid4())
    filename = req.url.split("/")[-1]
    if "?" in filename:
        filename = filename.split("?")[0]
    if not filename or "." not in filename:
        filename = "transfer_file.bin"
        
    tasks[task_id] = {
        "task_id": task_id,
        "url": req.url,
        "status": "downloading",
        "total_size": 0,
        "downloaded": 0,
        "speed": 0.0,
        "file_path": os.path.join(DATA_DIR, f"{task_id}_{filename}"),
        "original_filename": filename,
        "gofile_url": None,
        "timestamp": time.time()
    }
    background_tasks.add_task(process_gofile_transfer, task_id, req.url)
    return {"task_id": task_id}

# Chunked file reader generator to avoid streaming upload memory spikes
async def file_sender(file_path, chunk_size=256 * 1024):
    with open(file_path, 'rb') as f:
        while True:
            chunk = f.read(chunk_size)
            if not chunk:
                break
            yield chunk
            del chunk

async def process_gofile_transfer(task_id: str, url: str):
    file_path = tasks[task_id]["file_path"]
    try:
        # 1. Download Setup
        connector = aiohttp.TCPConnector(limit=10)
        async with aiohttp.ClientSession(connector=connector, read_bufsize=256 * 1024) as session:
            async with session.get(url, timeout=None) as response:
                response.raise_for_status()
                total_size = int(response.headers.get('Content-Length', 0))
                tasks[task_id]["total_size"] = total_size
                
                cd = response.headers.get('Content-Disposition')
                if cd and 'filename=' in cd:
                    fname = re.findall('filename="([^"]+)"', cd)
                    if not fname:
                        fname = re.findall('filename=([^;]+)', cd)
                    if fname:
                        new_filename = fname[0]
                        new_file_path = os.path.join(DATA_DIR, f"{task_id}_{new_filename}")
                        tasks[task_id]["file_path"] = new_file_path
                        tasks[task_id]["original_filename"] = new_filename
                        file_path = new_file_path

                downloaded = 0
                start_time = time.time()
                last_time = start_time
                last_downloaded = 0
                chunk_counter = 0
                
                with open(file_path, 'wb') as f:
                    async for chunk in response.content.iter_chunked(256 * 1024):
                        if not chunk:
                            break
                        f.write(chunk)
                        downloaded += len(chunk)
                        tasks[task_id]["downloaded"] = downloaded
                        
                        current_time = time.time()
                        if current_time - last_time >= 1.0:
                            tasks[task_id]["speed"] = (downloaded - last_downloaded) / (current_time - last_time)
                            last_time = current_time
                            last_downloaded = downloaded
                        
                        f.flush()
                        chunk_counter += 1
                        if chunk_counter % 20 == 0:
                            os.fsync(f.fileno())
                            del chunk
                            gc.collect()
        
        # 2. Upload to GoFile (STREAMING JUGAD WITH ZERO RAM IMPRINT)
        tasks[task_id]["status"] = "uploading_to_gofile"
        tasks[task_id]["speed"] = 0.0
        
        async with aiohttp.ClientSession() as session:
            async with session.get("https://api.gofile.io/servers") as resp:
                servers_data = await resp.json()
                server = servers_data["data"]["servers"][0]["name"]
            
            upload_url = f"https://{server}.gofile.io/uploadFile"
            
            # Form data wrapping without buffering entire file array into memory
            data = aiohttp.FormData()
            data.add_field('file', file_sender(file_path), filename=tasks[task_id]["original_filename"])
                
            async with session.post(upload_url, data=data) as upload_resp:
                upload_result = await upload_resp.json()
                if upload_result["status"] == "ok":
                    tasks[task_id]["gofile_url"] = upload_result["data"]["downloadPage"]
                    tasks[task_id]["status"] = "completed"
                else:
                    raise Exception("GoFile upload failed")
                        
        try:
            os.remove(file_path)
        except:
            pass
            
    except Exception as e:
        tasks[task_id]["status"] = "error"
        tasks[task_id]["error"] = str(e)
    finally:
        gc.collect()

# ==================== NEW ZIP MANAGEMENT FEATURES ====================

@app.get("/list_zip/{task_id}")
async def list_zip_contents(task_id: str):
    if task_id not in tasks:
        return {"error": "Task not found"}
    
    file_path = tasks[task_id]["file_path"]
    if not os.path.exists(file_path):
        return {"error": "File does not exist on server"}
        
    if not zipfile.is_zipfile(file_path):
        return {"error": "This file is not a valid ZIP archive"}
        
    try:
        with zipfile.ZipFile(file_path, 'r') as z:
            file_list = []
            for info in z.infolist():
                file_list.append({
                    "filename": info.filename,
                    "file_size_mb": round(info.file_size / (1024 * 1024), 2),
                    "is_dir": info.is_dir()
                })
            return {"task_id": task_id, "total_files": len(file_list), "files": file_list}
    except Exception as e:
        return {"error": f"Failed to read ZIP: {str(e)}"}

@app.post("/extract_zip/{task_id}")
async def extract_zip_file(task_id: str, background_tasks: BackgroundTasks):
    if task_id not in tasks:
        return {"error": "Task not found"}
        
    file_path = tasks[task_id]["file_path"]
    if not os.path.exists(file_path):
        return {"error": "File does not exist on server"}
        
    if not zipfile.is_zipfile(file_path):
        return {"error": "This file is not a valid ZIP archive"}
        
    extract_folder = os.path.join(DATA_DIR, f"extracted_{task_id}")
    os.makedirs(extract_folder, exist_ok=True)
    
    tasks[task_id]["status"] = "extracting"
    tasks[task_id]["extract_path"] = extract_folder
    
    background_tasks.add_task(process_zip_extraction, task_id, file_path, extract_folder)
    return {"status": "extraction_started", "task_id": task_id, "target_folder": extract_folder}

def process_zip_extraction(task_id: str, zip_path: str, target_dir: str):
    """Memory safe extraction processing member by member loop"""
    try:
        with zipfile.ZipFile(zip_path, 'r') as z:
            file_counter = 0
            for member in z.namelist():
                z.extract(member, path=target_dir)
                file_counter += 1
                # Release lock periodically inside loops
                if file_counter % 50 == 0:
                    gc.collect()
        tasks[task_id]["status"] = "extracted"
    except Exception as e:
        tasks[task_id]["status"] = "extraction_error"
        tasks[task_id]["error"] = str(e)
    finally:
        gc.collect()

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