# Author: Khadim Gueye
# ABCOMICS Genomic Data Upload Tool
# Two modes:
#   Manifest mode (new):  abcupload -d /path -m manifest.tsv [-k API_KEY]
#   Legacy mode (old):    abcupload -u USER -p PROJECT -d files/ [-k API_KEY]

from __future__ import annotations

import argparse
import csv
import glob
import os
import re
import sys
from pathlib import Path
from typing import Generator

import requests
from rich import box
from rich.console import Console
from rich.panel import Panel
from rich.progress import (
    BarColumn,
    DownloadColumn,
    Progress,
    TextColumn,
    TimeRemainingColumn,
    TransferSpeedColumn,
)
from rich.table import Table
from rich.text import Text

from abcupload import __version__

PROD_API = "https://abcomics-upload-492456311172.europe-west1.run.app"
TEST_API = "https://abcomics-test-upload-492456311172.europe-west1.run.app"

ALLOWED_SUFFIXES = {
    ".fastq", ".fq", ".fastq.gz", ".fq.gz", ".fastq.bz2", ".fq.bz2",
    ".fasta", ".fa", ".fna", ".faa", ".ffn", ".frn",
    ".fasta.gz", ".fa.gz", ".fna.gz", ".faa.gz",
    ".bam", ".sam", ".cram",
    ".bai", ".crai", ".csi", ".tbi",
    ".vcf", ".vcf.gz", ".bcf", ".bcf.gz",
    ".gff", ".gff3", ".gtf", ".bed",
    ".gff.gz", ".gff3.gz", ".gtf.gz", ".bed.gz",
    ".csv", ".tsv",
}

console     = Console()
err_console = Console(stderr=True)


# ── Helpers ───────────────────────────────────────────────────────────────────

def _has_bio_ext(name: str) -> bool:
    lower = name.lower()
    for ext in sorted(ALLOWED_SUFFIXES, key=len, reverse=True):
        if lower.endswith(ext):
            return True
    return False


def _human_size(n: int) -> str:
    for unit in ("B", "KB", "MB", "GB", "TB"):
        if n < 1024 or unit == "TB":
            return f"{n:.1f} {unit}" if unit != "B" else f"{n} B"
        n /= 1024  # type: ignore[assignment]
    return str(n)


# ── File resolution (legacy mode) ─────────────────────────────────────────────

def resolve_files(data_arg: str) -> list[Path]:
    """Return a list of Paths from a directory, glob, or comma-separated list."""
    path = Path(data_arg)

    if path.is_dir():
        files = sorted(f for f in path.iterdir() if f.is_file() and _has_bio_ext(f.name))
        if not files:
            console.print(f"  [yellow]![/] No bioinformatics files found in: [bold]{data_arg}[/]")
        return files

    results: list[Path] = []
    for part in data_arg.split(","):
        part = part.strip()
        if not part:
            continue
        matches = sorted(glob.glob(part))
        if matches:
            results.extend(Path(m) for m in matches if Path(m).is_file())
        elif Path(part).is_file():
            results.append(Path(part))
        else:
            console.print(f"  [yellow]![/] No match: [bold]{part}[/]")

    return results


# ── Manifest parsing ──────────────────────────────────────────────────────────

def parse_manifest(manifest_path: str) -> list[dict]:
    """
    Parse a TSV manifest file.
    Expected columns: entity_id  original_filename  new_gcs_filename  project_id
    Lines starting with # are treated as comments.
    """
    rows = []
    with open(manifest_path, newline="", encoding="utf-8") as fh:
        for lineno, raw in enumerate(fh, 1):
            line = raw.strip()
            if not line or line.startswith("#"):
                continue
            # Accept both tab and comma delimiters
            parts = line.split("\t") if "\t" in line else line.split(",")
            if len(parts) < 4:
                console.print(f"  [yellow]![/] Skipping line {lineno}: expected 4 columns, got {len(parts)}")
                continue
            entity_id, orig, new_name, project_id = (p.strip() for p in parts[:4])
            if not entity_id or not orig or not new_name or not project_id:
                console.print(f"  [yellow]![/] Skipping line {lineno}: empty column")
                continue
            rows.append({
                "entity_id":        entity_id,
                "original_filename": orig,
                "new_gcs_filename":  new_name,
                "project_id":        project_id,
            })
    return rows


# ── API call — get signed URL ─────────────────────────────────────────────────

def get_upload_url(
    api_url: str,
    api_key: str,
    project_id: str,
    filename: str,
    overwrite: bool = False,
    username: str = "",
) -> tuple[str, dict]:
    """
    POST to the API to get a GCS signed upload URL.
    Returns ('ok' | 'exists' | 'error', data_dict).
    """
    payload: dict = {"project_id": project_id, "filename": filename}
    if username:
        payload["username"] = username
    if overwrite:
        payload["overwrite"] = True

    try:
        resp = requests.post(
            f"{api_url}/generate-upload-url",
            json=payload,
            headers={"X-API-Key": api_key},
            timeout=30,
        )
    except requests.exceptions.ConnectionError:
        return "error", {"error": "Connection failed. Check your internet connection."}
    except requests.exceptions.Timeout:
        return "error", {"error": "Request timed out (30 s)."}
    except requests.exceptions.RequestException as exc:
        return "error", {"error": str(exc)}

    try:
        data = resp.json()
    except ValueError:
        data = {"error": f"Non-JSON response (HTTP {resp.status_code})"}

    if resp.status_code == 409:
        return "exists", data
    if resp.status_code != 200:
        return "error", data
    return "ok", data


# ── File upload — PUT to GCS ──────────────────────────────────────────────────

def upload_to_gcs(file_path: Path, upload_url: str) -> tuple[bool, str]:
    """Upload a file to GCS via a signed PUT URL with a rich progress bar."""
    file_size = file_path.stat().st_size

    with Progress(
        TextColumn("     "),
        BarColumn(bar_width=34),
        "[progress.percentage]{task.percentage:>3.0f}%",
        DownloadColumn(),
        TransferSpeedColumn(),
        TimeRemainingColumn(),
        console=console,
        transient=True,
    ) as progress:
        task = progress.add_task("upload", total=file_size)

        def _reader() -> Generator[bytes, None, None]:
            with open(file_path, "rb") as fh:
                while True:
                    chunk = fh.read(65536)
                    if not chunk:
                        break
                    progress.advance(task, len(chunk))
                    yield chunk

        try:
            resp = requests.put(
                upload_url,
                data=_reader(),
                headers={"Content-Type": "application/octet-stream"},
                timeout=None,
            )
        except requests.exceptions.RequestException as exc:
            return False, str(exc)

    if resp.status_code == 200:
        return True, ""
    return False, f"HTTP {resp.status_code}: {resp.text[:300]}"


# ── Per-file orchestration (manifest mode) ────────────────────────────────────

def upload_one_manifest(
    local_path: Path,
    row: dict,
    api_key: str,
    api_url: str,
    overwrite: bool,
    index: int,
    total: int,
) -> str:
    """Upload a single file using manifest metadata. Returns 'ok', 'fail', or 'skip'."""
    new_name = row["new_gcs_filename"]
    size_str = _human_size(local_path.stat().st_size)
    console.print(
        f"  [{index}/{total}] [bold]{local_path.name}[/] → [bold cyan]{new_name}[/]  [dim]{size_str}[/]"
    )

    status, data = get_upload_url(api_url, api_key, row["project_id"], new_name, overwrite)

    if status == "exists":
        console.print("     [yellow]![/] File already exists in the project directory.")
        try:
            answer = input("     Overwrite? [y/N] ").strip().lower()
        except (EOFError, KeyboardInterrupt):
            answer = "n"
        if answer == "y":
            status, data = get_upload_url(api_url, api_key, row["project_id"], new_name, overwrite=True)
        else:
            console.print("     [dim]−[/] Skipped.")
            return "skip"

    if status != "ok":
        console.print(f"     [red]✗[/] {data.get('error', 'Unknown error')}")
        return "fail"

    upload_url = data.get("upload_url", "")
    if not upload_url:
        console.print("     [red]✗[/] Could not retrieve upload URL.")
        return "fail"

    ok, err_msg = upload_to_gcs(local_path, upload_url)

    if ok:
        console.print("     [green]✓[/] Uploaded successfully")
        return "ok"

    console.print(f"     [red]✗[/] Upload failed: {err_msg}")
    return "fail"


# ── Per-file orchestration (legacy mode) ──────────────────────────────────────

def upload_one(
    file_path: Path,
    username: str,
    project_id: str,
    api_key: str,
    test_mode: bool,
    overwrite: bool,
    index: int,
    total: int,
) -> str:
    """Upload a single file (legacy mode). Returns 'ok', 'fail', or 'skip'."""
    api_url  = TEST_API if test_mode else PROD_API
    filename = file_path.name
    size_str = _human_size(file_path.stat().st_size)

    console.print(f"  [{index}/{total}] [bold]{filename}[/]  [dim]{size_str}[/]")

    status, data = get_upload_url(api_url, api_key, project_id, filename, overwrite, username)

    if status == "exists":
        console.print("     [yellow]![/] File already exists in your project directory.")
        try:
            answer = input("     Overwrite? [y/N] ").strip().lower()
        except (EOFError, KeyboardInterrupt):
            answer = "n"
        if answer == "y":
            status, data = get_upload_url(
                api_url, api_key, project_id, filename, overwrite=True, username=username
            )
        else:
            console.print("     [dim]−[/] Skipped.")
            return "skip"

    if status != "ok":
        err_msg = data.get("error", "Unknown error")
        console.print(f"     [red]✗[/] {err_msg}")
        return "fail"

    upload_url = data.get("upload_url", "")
    if not upload_url:
        console.print("     [red]✗[/] Could not retrieve upload URL.")
        return "fail"

    ok, err_msg = upload_to_gcs(file_path, upload_url)

    if ok:
        console.print("     [green]✓[/] Uploaded successfully")
        return "ok"

    console.print(f"     [red]✗[/] Upload failed: {err_msg}")
    return "fail"


# ── Display helpers ───────────────────────────────────────────────────────────

def print_banner_manifest(directory: str, manifest: str) -> None:
    grid = Table.grid(padding=(0, 2))
    grid.add_column(style="dim")
    grid.add_column()
    grid.add_row("Directory", f"[bold]{directory}[/]")
    grid.add_row("Manifest",  f"[bold]{manifest}[/]")
    console.print(Panel(
        grid,
        title=f"[bold blue]ABCOMICS[/] Genomic Upload  [dim]v{__version__}[/]",
        border_style="blue",
        padding=(0, 1),
    ))
    console.print()


def print_banner(username: str, project_id: str, test_mode: bool) -> None:
    mode = (
        Text("TEST  ·  files deleted after 24 h", style="bold yellow")
        if test_mode
        else Text("Production", style="bold green")
    )
    grid = Table.grid(padding=(0, 2))
    grid.add_column(style="dim")
    grid.add_column()
    grid.add_row("User",    f"[bold]{username}[/]")
    grid.add_row("Project", f"[bold]{project_id}[/]")
    grid.add_row("Mode",    mode)
    console.print(Panel(
        grid,
        title=f"[bold blue]ABCOMICS[/] Genomic Upload  [dim]v{__version__}[/]",
        border_style="blue",
        padding=(0, 1),
    ))
    console.print()


def print_summary(results: dict) -> None:
    total = sum(results.values())
    tbl = Table(box=box.SIMPLE, show_header=False, padding=(0, 1))
    tbl.add_column(style="dim", width=10)
    tbl.add_column()
    tbl.add_row("Total", f"{total} file(s)")
    if results["ok"]:
        tbl.add_row("Uploaded", f"[green]{results['ok']}[/]")
    if results["fail"]:
        tbl.add_row("Failed",   f"[red]{results['fail']}[/]")
    if results["skip"]:
        tbl.add_row("Skipped",  f"[dim]{results['skip']}[/]")
    console.print()
    console.rule(style="dim")
    console.print(tbl)
    console.rule(style="dim")
    console.print()

    if results["fail"]:
        console.print("[red]Some files failed to upload. Check errors above.[/]\n")
    elif results["ok"]:
        console.print("[green]All files uploaded successfully.[/]\n")
    else:
        console.print("[yellow]No files were uploaded.[/]\n")


# ── Argument parser ───────────────────────────────────────────────────────────

def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="abcupload",
        description="ABCOMICS Genomic Data Upload Tool",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""\
manifest mode (recommended):
  abcupload -d /path/to/files -m PRJAB00001_upload_manifest.tsv
  abcupload -d /path/to/files -m manifest.tsv -k YOUR_API_KEY

  The manifest TSV is downloaded from the ABCOMICS dashboard after submitting
  runs and analyses metadata. It maps each original filename to its ID-based
  GCS destination name and project ID.

legacy mode:
  abcupload -u abc-000001 -p PRJAB00001 -d sample_R1.fastq.gz
  abcupload -u abc-000001 -p PRJAB00001 -d '*.fastq.gz'
  abcupload -u abc-000001 -p PRJAB00001 -d /path/to/data/ -t

environment variable:
  ABCOMICS_API_KEY   your API key (alternative to -k)
  export ABCOMICS_API_KEY=your_key_here
""",
        add_help=False,
    )

    req = parser.add_argument_group("manifest mode (recommended)")
    req.add_argument(
        "-m", dest="manifest", metavar="MANIFEST_TSV", default=None,
        help="TSV manifest downloaded from the ABCOMICS dashboard",
    )
    req.add_argument(
        "-d", dest="data", metavar="DIRECTORY_OR_FILES", required=True,
        help="directory containing files (manifest mode) or file / glob / comma-list (legacy mode)",
    )

    leg = parser.add_argument_group("legacy mode")
    leg.add_argument("-u", dest="username",   metavar="USERNAME",   default=None, help="ABCOMICS username")
    leg.add_argument("-p", dest="project_id", metavar="PROJECT_ID", default=None, help="project ID (e.g. PRJAB00001)")

    opt = parser.add_argument_group("options")
    opt.add_argument(
        "-k", dest="api_key", metavar="API_KEY",
        default=os.environ.get("ABCOMICS_API_KEY", ""),
        help="API key  (or: export ABCOMICS_API_KEY=...)",
    )
    opt.add_argument("-t", dest="test",      action="store_true", help="test mode — files deleted after 24 h")
    opt.add_argument("-o", dest="overwrite", action="store_true", help="overwrite existing files without asking")
    opt.add_argument("-h", "--help",    action="help",    help="show this help and exit")
    opt.add_argument("-V", "--version", action="version", version=f"%(prog)s {__version__}")

    return parser


# ── Entry point ───────────────────────────────────────────────────────────────

def main() -> None:
    parser = build_parser()
    args   = parser.parse_args()

    if not args.api_key:
        err_console.print(
            "[red]Error:[/] API key required.\n"
            "  Use [bold]-k KEY[/] or set [bold]ABCOMICS_API_KEY[/] environment variable.\n"
            "  Contact [bold]contact@abcomics.org[/] if you have not received your key."
        )
        sys.exit(1)

    # ── Manifest mode ─────────────────────────────────────────────────────────
    if args.manifest:
        manifest_path = args.manifest
        if not Path(manifest_path).is_file():
            err_console.print(f"[red]Error:[/] Manifest file not found: [bold]{manifest_path}[/]")
            sys.exit(1)

        directory = args.data
        if not Path(directory).is_dir():
            err_console.print(f"[red]Error:[/] Directory not found: [bold]{directory}[/]")
            sys.exit(1)

        rows = parse_manifest(manifest_path)
        if not rows:
            err_console.print("[red]Error:[/] No valid rows found in manifest.")
            sys.exit(1)

        api_url = TEST_API if args.test else PROD_API
        print_banner_manifest(directory, manifest_path)
        console.print(f"  Found [bold]{len(rows)}[/] file(s) in manifest:\n")

        results = {"ok": 0, "fail": 0, "skip": 0}

        for i, row in enumerate(rows, 1):
            local_path = Path(directory) / row["original_filename"]
            if not local_path.exists():
                console.print(
                    f"  [{i}/{len(rows)}] [yellow]![/] File not found locally: "
                    f"[bold]{row['original_filename']}[/] — skipping."
                )
                results["skip"] += 1
                console.print()
                continue

            outcome = upload_one_manifest(
                local_path = local_path,
                row        = row,
                api_key    = args.api_key,
                api_url    = api_url,
                overwrite  = args.overwrite,
                index      = i,
                total      = len(rows),
            )
            results[outcome] += 1
            console.print()

        print_summary(results)
        sys.exit(1 if results["fail"] > 0 else 0)

    # ── Legacy mode ───────────────────────────────────────────────────────────
    if not args.username or not args.project_id:
        err_console.print(
            "[red]Error:[/] Provide either:\n"
            "  Manifest mode: [bold]-d DIRECTORY -m MANIFEST_TSV[/]\n"
            "  Legacy mode:   [bold]-u USERNAME -p PROJECT_ID -d FILES[/]"
        )
        sys.exit(1)

    if not re.match(r"^PRJAB\d+$", args.project_id):
        err_console.print(
            "[red]Error:[/] PROJECT_ID must start with [bold]PRJAB[/] followed by digits "
            "(e.g. [bold]PRJAB00001[/])."
        )
        sys.exit(1)

    files = resolve_files(args.data)
    if not files:
        err_console.print(f"[red]Error:[/] No valid files found in: [bold]{args.data}[/]")
        sys.exit(1)

    api_url = TEST_API if args.test else PROD_API
    print_banner(args.username, args.project_id, args.test)
    console.print(f"  Found [bold]{len(files)}[/] file(s) to upload:\n")

    results = {"ok": 0, "fail": 0, "skip": 0}

    for i, fp in enumerate(files, 1):
        outcome = upload_one(
            file_path  = fp,
            username   = args.username,
            project_id = args.project_id,
            api_key    = args.api_key,
            test_mode  = args.test,
            overwrite  = args.overwrite,
            index      = i,
            total      = len(files),
        )
        results[outcome] += 1
        console.print()

    print_summary(results)
    sys.exit(1 if results["fail"] > 0 else 0)