"""Fetch the essay datasets from Kaggle straight into the Modal volume `essay-data` (nothing lands on the laptop).

    MODAL_PROFILE=labs3 modal run scripts/get_data.py
"""
import os

import modal

app = modal.App("essay-coach-get-data")
vol = modal.Volume.from_name("essay-data", create_if_missing=True)
image = modal.Image.debian_slim(python_version="3.11").pip_install("kaggle==1.6.17", "requests").apt_install("unzip")
token = open(os.path.expanduser("~/.kaggle/token")).read().strip() if modal.is_local() else ""

COMPETITIONS = ["feedback-prize-english-language-learning", "learning-agency-lab-automated-essay-scoring-2",
                "feedback-prize-2021", "feedback-prize-effectiveness"]
DATASETS = ["nbroad/persaude-corpus-2", "mazlumi/ielts-writing-scored-essays-dataset", "xntrng15/ielts-writing-dataset"]
GITHUB = ["scrosseye/ELLIPSE-Corpus", "scrosseye/persuade_corpus_2.0"]  # every csv/xlsx in the repo root


@app.function(image=image, volumes={"/data": vol}, timeout=3600, secrets=[modal.Secret.from_dict({"KAGGLE_API_TOKEN": token})])
def fetch():
    import subprocess
    out = {}
    for c in COMPETITIONS:
        d = f"/data/{c}"
        if os.path.isdir(d) and os.listdir(d):
            out[c] = "already there"
            continue
        os.makedirs(d, exist_ok=True)
        r = subprocess.run(["kaggle", "competitions", "download", "-c", c, "-p", d], capture_output=True, text=True)
        if r.returncode:
            out[c] = "FAILED: " + (r.stderr or r.stdout)[-300:]
            continue
        subprocess.run(f"cd {d} && unzip -oq '*.zip' && rm -f *.zip", shell=True, check=True)
        out[c] = sorted(os.listdir(d))[:8]
    for ds in DATASETS:
        d = f"/data/{ds.split('/')[-1]}"
        if os.path.isdir(d) and os.listdir(d):
            out[ds] = "already there"
            continue
        os.makedirs(d, exist_ok=True)
        r = subprocess.run(["kaggle", "datasets", "download", "-d", ds, "-p", d, "--unzip"], capture_output=True, text=True)
        out[ds] = sorted(os.listdir(d))[:8] if not r.returncode else "FAILED: " + (r.stderr or r.stdout)[-300:]
    import requests
    for repo in GITHUB:
        d = f"/data/{repo.split('/')[-1]}"
        os.makedirs(d, exist_ok=True)
        got = []
        for item in requests.get(f"https://api.github.com/repos/{repo}/contents", timeout=60).json():
            name = item.get("name", "")
            if name.lower().endswith((".csv", ".xlsx", ".md", ".txt", ".zip")) and item.get("download_url"):
                with open(f"{d}/{name}", "wb") as f:
                    f.write(requests.get(item["download_url"], timeout=600).content)
                got.append(name)
        out[repo] = got
    vol.commit()
    return out


@app.local_entrypoint()
def main():
    for k, v in fetch.remote().items():
        print(k, v)