#!/usr/bin/env python3

import os
import datetime
import requests
import gzip
from pathlib import Path

# Define resource and paths
resource = "ensembl"
base_dir = Path(f"/data/projects/glygen/downloads/{resource}")
shared_dir = Path(f"/data/shared/glygen/downloads/{resource}")

# Create timestamped folder
new_dir = datetime.datetime.now().strftime("%Y_%m_%d")
target_dir = base_dir / new_dir
target_dir.mkdir(parents=True, exist_ok=True)

# URLs and update TSV filenames
files_to_download = {
    "https://ftp.ensembl.org/pub/release-115/gtf/drosophila_melanogaster/Drosophila_melanogaster.BDGP6.54.115.chr.gtf.gz": "fruitfly_ensembl_coords.tsv",
    "https://ftp.ensemblgenomes.ebi.ac.uk/pub/plants/release-62/gff3/arabidopsis_thaliana/Arabidopsis_thaliana.TAIR10.62.gff3.gz": "arabidopsis_ensembl_coords.tsv",
    "https://ftp.ensemblgenomes.ebi.ac.uk/pub/fungi/release-62/gtf/saccharomyces_cerevisiae/Saccharomyces_cerevisiae.R64-1-1.62.gtf.gz": "yeast_ensembl_coords.tsv"
}

headers = {"User-Agent": "Mozilla"}
column_header = "\t".join(["seqid", "source", "type", "start", "end", "score", "strand", "phase", "attributes"])

# Download, unzip, clean, and convert to TSV
for url, final_tsv in files_to_download.items():
    gz_name = os.path.basename(url)
    gz_path = target_dir / gz_name
    txt_path = target_dir / gz_name.replace(".gz", "")
    tsv_path = target_dir / final_tsv

    print(f"Downloading {gz_name} ...")
    response = requests.get(url, headers=headers)
    response.raise_for_status()
    with open(gz_path, "wb") as f:
        f.write(response.content)

    print(f"Unzipping {gz_name} ...")
    with gzip.open(gz_path, "rt") as f_in, open(txt_path, "w") as f_out:
        for line in f_in:
            if not line.startswith("#"):
                f_out.write(line)
    gz_path.unlink()

    print(f"Converting and adding headers to {final_tsv} ...")
    with open(txt_path, "r") as fin, open(tsv_path, "w") as fout:
        fout.write(column_header + "\n")
        for line in fin:
            # GFF3/GTF files already use tabs - just write as-is
            fout.write(line)
    txt_path.unlink()

# Update permissions
os.chdir(shared_dir.parent)
os.chmod(shared_dir / new_dir, 0o775)
for root, dirs, files in os.walk(shared_dir / new_dir):
    for d in dirs:
        os.chmod(os.path.join(root, d), 0o775)
    for f in files:
        os.chmod(os.path.join(root, f), 0o775)

# Update symbolic link
os.chdir(base_dir)
link_path = base_dir / "current"
if link_path.exists() or link_path.is_symlink():
    link_path.unlink()
os.symlink(new_dir, link_path)

print("Download, cleanup, TSV conversion, and header addition complete.")