import os
import time
import requests
from urllib.parse import urlparse
from huggingface_hub import upload_file
from fastapi import FastAPI
from contextlib import asynccontextmanager
import asyncio
import logging

# === CONFIGURATION ===
HF_TOKEN = os.environ.get("HF_TOKEN")
REPO_ID = "Fred808/BG1"
DATA_PATH = "DaVinci"
OUTPUT_DIR = "batch_downloads"
DOWNLOAD_URLS = [
    "https://ww6.zeroupload.xyz/8931dc42eee2300803892aebe177286f/FF_DavinciResolveEditingWorkflow_DownloadPirate.com.rar?download_token=96d34e6de5ef7f0441d4c73619580dd419740fe1e5981da9c44b19b59ea83a4f",
    "https://ww6.zeroupload.xyz/795910b2a995ae221bde85bf8a36f920/GCC_MasterTrainingDaVinciResolve_DownloadPirate.com.rar?download_token=d4f96f8a6414a89310f266d643cd4d88e688452b6ed69b29d5e423614f7c69ef",
    "https://ww6.zeroupload.xyz/9c024aeb4ba28cd22201357936fe95af/GC_ProColorGradingDaVinciResolve_DownloadPirate.com.rar?download_token=8bf4a9da56885c9fd615b822596fbb9afd0ef548b3283895e9f35b540197dd56",
    
]
DELAY_BETWEEN_DOWNLOADS = 30  # seconds

# === Setup Logging ===
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")

# === Prepare output folder ===
os.makedirs(OUTPUT_DIR, exist_ok=True)

app = FastAPI()

# === DUMMY ROUTE TO KEEP SERVER HEALTHY ===
@app.get("/")
def keep_alive():
    return {"status": "running"}

# === Upload Function ===
def upload_to_dataset(filepath):
    try:
        upload_file(
            path_or_fileobj=filepath,
            path_in_repo=f"{DATA_PATH}/{os.path.basename(filepath)}",
            repo_id=REPO_ID,
            repo_type="dataset",
            token=HF_TOKEN
        )
        logging.info(f"[↑] Uploaded: {filepath}")
    except Exception as e:
        logging.error(f"[!] Upload failed: {filepath} — {e}")

# === Background Worker ===
async def downloader_worker():
    for direct_download_link in DOWNLOAD_URLS:
        logging.info("[*] Waiting before next download...")
        await asyncio.sleep(DELAY_BETWEEN_DOWNLOADS)

        try:
            logging.info(f"[*] Downloading from: {direct_download_link}")
            filename = os.path.basename(urlparse(direct_download_link).path)
            if not filename or "." not in filename:
                filename = "downloaded_file_" + str(int(time.time()))

            local_path = os.path.join(OUTPUT_DIR, filename)
            logging.info(f"[*] Saving to: {local_path}")

            with requests.get(direct_download_link, stream=True) as r:
                r.raise_for_status()
                with open(local_path, "wb") as f:
                    for chunk in r.iter_content(chunk_size=8192):
                        f.write(chunk)

            logging.info(f"[✓] Downloaded: {filename}")
            upload_to_dataset(local_path)
            os.remove(local_path)

        except Exception as e:
            logging.error(f"[!] Error with {direct_download_link}: {e}")

    logging.info("✅ All files processed.")

@app.get("/")
def stay_alive():
    return {"msg": "Running"}

@app.get("/health")
def healthcheck():
    return {"healthy": True}

# === FastAPI Lifespan ===
@asynccontextmanager
async def lifespan(app: FastAPI):
    logging.info("🚀 Starting FastAPI download-uploader microservice...")
    task = asyncio.create_task(downloader_worker())
    yield
    task.cancel()
    logging.info("🛑 Shutting down microservice.")

# === FastAPI App ===
app = FastAPI(lifespan=lifespan)
# Re-assign app with lifespan logic
app.router.lifespan_context = lifespan