# Author: Khadim Gueye
# ABCOMICS Genomic Data Upload Tool
# Handles: URL signing (POST) + file upload (PUT), progress, duplicate detection, test mode.

from __future__ import annotations

import argparse
import glob
import os
import re
import sys
from pathlib import Path
from typing import Generator, Optional

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 ───────────────────────────────────────────────────────────

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


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

def get_upload_url(
    api_url: str,
    api_key: str,
    username: str,
    project_id: str,
    filename: str,
    overwrite: bool = False,
) -> tuple[str, dict]:
    """POST to the API to get a GCS signed upload URL.

    Returns (status, data) where status is 'ok', 'exists', or 'error'.
    """
    payload: dict = {"username": username, "project_id": project_id, "filename": filename}
    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",
                    "Content-Length": str(file_size),
                },
                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}"


# ── Per-file orchestration ────────────────────────────────────────────────────

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. 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, username, project_id, filename, overwrite)

    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, username, project_id, filename, overwrite=True
            )
        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")
        if test_mode and "auto_delete" in data:
            console.print(f"     [dim]{data['auto_delete']}[/]")
        return "ok"

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


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

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="""\
examples:
  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 file1.bam,file2.bam -o
  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("required")
    req.add_argument(
        "-u", dest="username", metavar="USERNAME", required=True,
        help="ABCOMICS username",
    )
    req.add_argument(
        "-p", dest="project_id", metavar="PROJECT_ID", required=True,
        help="project ID  (e.g. PRJAB00001)",
    )
    req.add_argument(
        "-d", dest="data", metavar="DATA", required=True,
        help="file(s): single file | comma-list | glob pattern | directory",
    )

    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)

    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)

    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)