"""Fetch the PERSUADE 2.0 corpus (argument parts with effectiveness) from its Google Drive links into the volume."""
import modal

app = modal.App("essay-coach-get-persuade")
vol = modal.Volume.from_name("essay-data")
image = modal.Image.debian_slim(python_version="3.11").pip_install("gdown", "pandas").apt_install("p7zip-full")
FILES = {"persuade_2.0_train.csv": "13phHyDzIsb0MHyJr6q-B-qIa9P2tM135", "persuade_2.0_test.zip": "1K1SIJiG-2zWgMlTzxQeYOcLwOsFaVel1"}


@app.function(image=image, volumes={"/data": vol}, timeout=3600)
def fetch():
    import os
    import subprocess
    import pandas as pd
    d = "/data/persuade_corpus_2.0"
    os.makedirs(d, exist_ok=True)
    out = {}
    for name, fid in FILES.items():
        p = f"{d}/{name}"
        if not os.path.exists(p):
            r = subprocess.run(["gdown", fid, "-O", p], capture_output=True, text=True)
            if r.returncode or not os.path.exists(p):
                out[name] = "FAILED: " + (r.stderr or r.stdout)[-300:]
                continue
        out[name] = f"{os.path.getsize(p)/1e6:.0f} MB"
        if name.endswith(".zip"):
            subprocess.run(["7z", "x", "-y", "-ppersuade_test", f"-o{d}", p], capture_output=True, text=True)
            out[name] += " -> " + str([f for f in os.listdir(d) if "test" in f])
    for f in os.listdir(d):
        if f.endswith(".csv"):
            df = pd.read_csv(f"{d}/{f}", nrows=3)
            out[f + " cols"] = list(df.columns)
    vol.commit()
    return out


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