app.py

import os
import urllib.parse
from fastapi import FastAPI, HTTPException, Request, status
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse, FileResponse
from pydantic import BaseModel

app = FastAPI()

# Configuração de CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

STORAGE_DIR = "/data"

if not os.path.exists(STORAGE_DIR):
    os.makedirs(STORAGE_DIR, exist_ok=True)

class UploadURLRequest(BaseModel):
    fileName: str
    contentType: str

class DeleteRequest(BaseModel):
    key: str

@app.get("/")
def read_root():
    return {"ok": True, "message": "Hugging Face storage node is running!", "storage": STORAGE_DIR}

@app.get("/health")
def health():
    return {"ok": True, "service": "huggingface-local-bucket"}

@app.get("/files")
def list_files():
    try:
        files = []
        for entry in os.scandir(STORAGE_DIR):
            if entry.is_file():
                if entry.name.startswith("."):
                    continue
                stat = entry.stat()
                files.append({
                    "fileName": entry.name,
                    "size": stat.st_size,
                    "lastModified": stat.st_mtime * 1000
                })
        files.sort(key=lambda x: x["lastModified"], reverse=True)
        return {"ok": True, "files": files, "bucket": "HuggingFace (Mount)"}
    except Exception as e:
        return {"ok": False, "error": str(e)}

@app.post("/upload-url")
def upload_url(payload: UploadURLRequest, request: Request):
    try:
        filename_encoded = urllib.parse.quote(payload.fileName)
        host_header = request.headers.get("host", "")
        if "hf.space" in host_header:
            base_url = f"https://{host_header}/"
        else:
            base_url = str(request.base_url)
            if base_url.startswith("http://") and "localhost" not in base_url and "127.0.0.1" not in base_url:
                base_url = base_url.replace("http://", "https://", 1)
        return {"ok": True, "url": f"{base_url}upload-file/{filename_encoded}"}
    except Exception as e:
        return {"ok": False, "error": str(e)}

@app.put("/upload-file/{filename}")
async def receive_upload(filename: str, request: Request):
    try:
        filename_decoded = urllib.parse.unquote(filename)
        filepath = os.path.join(STORAGE_DIR, filename_decoded)
        if not os.path.abspath(filepath).startswith(os.path.abspath(STORAGE_DIR)):
            raise HTTPException(status_code=400, detail="Caminho inválido")
        with open(filepath, "wb") as f:
            async for chunk in request.stream():
                f.write(chunk)
        return {"ok": True}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# 🎬 NOVA ROTA DE DOWNLOAD COM SUPORTE A STREAMING REAL (RANGE REQUESTS)
@app.get("/download")
def download_file(name: str, request: Request):
    try:
        filepath = os.path.join(STORAGE_DIR, name)
        if not os.path.exists(filepath):
            raise HTTPException(status_code=404, detail="Arquivo não encontrado")
        
        file_size = os.path.getsize(filepath)
        range_header = request.headers.get("range", "")
        
        # Define o tipo de mídia adequado
        lower_name = name.lower()
        media_type = "video/mp4" if lower_name.endswith(".mp4") else "application/octet-stream"
        if lower_name.endswith(".webm"): media_type = "video/webm"
        elif lower_name.endswith(".mkv"): media_type = "video/x-matroska"
        elif lower_name.endswith(".mp3"): media_type = "audio/mpeg"

        if range_header and range_header.startswith("bytes="):
            try:
                range_value = range_header.replace("bytes=", "").strip()
                if "-" in range_value:
                    start_str, end_str = range_value.split("-", 1)
                    start = int(start_str) if start_str else 0
                    end = int(end_str) if end_str else file_size - 1
                else:
                    start = int(range_value)
                    end = file_size - 1
            except ValueError:
                start = 0
                end = file_size - 1

            start = max(0, min(start, file_size - 1))
            end = max(start, min(end, file_size - 1))
            content_length = (end - start) + 1

            def file_iterator():
                with open(filepath, "rb") as f:
                    f.seek(start)
                    remaining = content_length
                    while remaining > 0:
                        chunk_size = min(4096 * 16, remaining)  # 64KB chunks
                        chunk = f.read(chunk_size)
                        if not chunk:
                            break
                        remaining -= len(chunk)
                        yield chunk

            headers = {
                "Content-Range": f"bytes {start}-{end}/{file_size}",
                "Accept-Ranges": "bytes",
                "Content-Length": str(content_length),
                "Cache-Control": "public, max-age=2592000",
            }
            return StreamingResponse(file_iterator(), status_code=206, headers=headers, media_type=media_type)
        
        # Caso não haja Range Request, envia normal equilibrando o cache
        response = FileResponse(filepath, media_type=media_type, filename=name)
        response.headers["Accept-Ranges"] = "bytes"
        return response
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/delete")
def delete_file(payload: DeleteRequest):
    try:
        filepath = os.path.join(STORAGE_DIR, payload.key)
        if os.path.exists(filepath):
            os.remove(filepath)
            return {"ok": True}
        return {"ok": False, "error": "Arquivo não encontrado"}
    except Exception as e:
        return {"ok": False, "error": str(e)}

@app.get("/preview/{filename:path}")
def preview_text(filename: str):
    try:
        filepath = os.path.join(STORAGE_DIR, filename)
        if not os.path.exists(filepath):
            return {"ok": False, "error": "Arquivo não encontrado"}
        with open(filepath, "rb") as f:
            content = f.read()
        try:
            text = content.decode('utf-8')
        except Exception:
            text = content.decode('latin-1')
        return {"ok": True, "content": text, "filename": filename}
    except Exception as e:
        return {"ok": False, "error": str(e)}