#!/usr/bin/env python3
"""Warhammer DLP Helper.

Give it a Discord channel/thread id and an output directory; it exports the
channel, extracts the encoded download links (gofile.io and 0807.st), picks
the right downloader for each link, and keeps going when individual links
fail. Everything is recorded in state.json; failures also land in
failed_urls.txt so a re-run (or --urls-file) retries exactly those.

Usage:
    python3 dlp.py <channel_id> <out_dir> [options]
"""

from __future__ import annotations

import argparse
import getpass
import hashlib
import json
import logging
import os
import sys
import time
import zipfile
from pathlib import Path

from dlphelper import discord as dlp_discord
from dlphelper import gofile as dlph_gofile
from dlphelper import zero807 as dlph_zero807
from dlphelper.state import State, append_manifest

log = logging.getLogger("dlp")


def sha256_of(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as fh:
        for chunk in iter(lambda: fh.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def resolve_token(cli_value: str | None) -> str:
    if cli_value:
        return cli_value
    if os.environ.get("DISCORD_TOKEN"):
        return os.environ["DISCORD_TOKEN"]
    log.info("No Discord token given ($DISCORD_TOKEN); prompting (input hidden).")
    token = getpass.getpass("Discord token: ")
    if not token:
        raise SystemExit("No token provided, aborting.")
    return token


def load_urls_file(path: Path) -> list[dlp_discord.FoundUrl]:
    """Load a plain URL list (e.g. failed_urls.txt); no attribution available."""
    found: dict[str, dlp_discord.FoundUrl] = {}
    for line in path.read_text(encoding="utf-8").splitlines():
        url = line.strip()
        if not url or url.startswith("#"):
            continue
        host = dlp_discord.classify(url)
        if host is None:
            log.warning("Unsupported host in %s: %s", path.name, url)
            continue
        found.setdefault(url, dlp_discord.FoundUrl(url=url, host=host))
    return sorted(found.values(), key=lambda f: f.url)


def record_files(state: State, url: str, files: list[dict], source: dict) -> None:
    state.mark(url, status="ok", error=None, files=files)
    for f in files:
        append_manifest(Path(source["manifest"]), {
            "channel": source["channel"],
            "url": url,
            "host": source["host"],
            "message_id": source.get("message_id"),
            "timestamp": source.get("timestamp"),
            "author": source.get("author"),
            **f,
        })


def extract_zips(files: list[dict]) -> None:
    """Optional: unzip each downloaded archive next to itself."""
    for f in files:
        path = Path(f["path"])
        if path.suffix.lower() != ".zip" or not path.exists():
            continue
        target = path.with_suffix("")
        if target.exists() and any(target.iterdir()):
            log.info("Skip extracting %s (target exists)", path.name)
            continue
        log.info("Extracting %s -> %s", path.name, target.name)
        try:
            target.mkdir(parents=True, exist_ok=True)
            with zipfile.ZipFile(path) as zf:
                zf.extractall(target)
        except (zipfile.BadZipFile, OSError) as exc:
            log.warning("Extraction of %s failed: %s", path.name, exc)


def run_gofile(
    urls: list[dlp_discord.FoundUrl],
    args: argparse.Namespace,
    state: State,
    source_by_url: dict[str, dict],
) -> None:
    client = dlph_gofile.GoFileClient(password=args.password)
    try:
        client.ensure_token()
    except dlph_gofile.GoFileError as exc:
        log.error("%s - marking all pending gofile links as failed.", exc)
        for f in urls:
            if state.is_done(f.url):
                continue
            state.mark(f.url, status="failed", error=str(exc),
                       files=[])
        return

    consecutive_failures = 0
    for f in urls:
        if state.is_done(f.url):
            continue
        source = source_by_url[f.url]

        for attempt in range(1, args.retries + 1):
            try:
                files_spec = client.list_files(f.url)
                if not files_spec:
                    raise dlph_gofile.GoFileError("no files found under URL")
                files: list[dict] = []
                for relpath, name, link in files_spec:
                    dest = Path(args.out_dir) / "downloads" / "gofile" / relpath / name
                    size = client.download_file(link, dest)
                    entry = {"path": str(dest), "size": size, "name": name}
                    if args.hash:
                        entry["sha256"] = sha256_of(dest)
                    files.append(entry)
                log.info("OK %s (%d files)", f.url, len(files))
                record_files(state, f.url, files, source)
                if args.extract:
                    extract_zips(files)
                consecutive_failures = 0
                break
            except dlph_gofile.GoFilePasswordRequired as exc:
                # Retrying without a password never helps
                log.error("PASSWORD REQUIRED %s: %s", f.url, exc)
                state.mark(f.url, status="failed", error=str(exc), files=[])
                break
            except dlph_gofile.GoFileAuthError as exc:
                # Token may have died mid-run (rate limit); re-auth once per attempt
                log.warning("Auth error for %s: %s - re-authenticating", f.url, exc)
                state.mark(f.url, error=str(exc), files=[])
                client.reset()
                try:
                    client.ensure_token()
                except dlph_gofile.GoFileError as reauth_exc:
                    log.error("Re-authentication failed: %s", reauth_exc)
                if attempt < args.retries:
                    time.sleep(args.retry_delay * attempt)
            except dlph_gofile.GoFileError as exc:
                log.warning("Attempt %d/%d failed for %s: %s",
                            attempt, args.retries, f.url, exc)
                state.mark(f.url, error=str(exc), files=[])
                if attempt < args.retries:
                    time.sleep(args.retry_delay * attempt)
            except Exception as exc:  # noqa: BLE001 - a bug must not stop the run
                log.exception("Unexpected error for %s: %s", f.url, exc)
                state.mark(f.url, error=f"unexpected: {exc}", files=[])
                if attempt < args.retries:
                    time.sleep(args.retry_delay * attempt)
            except KeyboardInterrupt:
                state.mark(f.url, error="interrupted", files=[])
                raise
        else:
            rec = state.records.get(f.url, {})
            state.mark(f.url, status="failed", error=rec.get("error", "unknown"), files=[])
            consecutive_failures += 1
            log.error("FAILED %s after %d attempts", f.url, args.retries)
            if consecutive_failures >= args.cooldown_failures:
                log.warning(
                    "%d failures in a row - cooling down %ds before continuing "
                    "(rate limit?). The run does not stop.",
                    consecutive_failures, args.cooldown_seconds,
                )
                time.sleep(args.cooldown_seconds)


def run_0807(
    urls: list[dlp_discord.FoundUrl],
    args: argparse.Namespace,
    state: State,
    source_by_url: dict[str, dict],
) -> None:
    out_dir = Path(args.out_dir) / "downloads" / "0807"
    consecutive_failures = 0
    for f in urls:
        if state.is_done(f.url):
            continue
        source = source_by_url[f.url]

        for attempt in range(1, args.retries + 1):
            try:
                dest, size = dlph_zero807.download(f.url, out_dir)
                entry: dict = {"path": str(dest), "size": size, "name": dest.name}
                if args.hash:
                    entry["sha256"] = sha256_of(dest)
                log.info("OK %s (%s)", f.url, format_size(size))
                record_files(state, f.url, [entry], source)
                if args.extract:
                    extract_zips([entry])
                consecutive_failures = 0
                break
            except dlph_zero807.Zero807Error as exc:
                log.warning("Attempt %d/%d failed for %s: %s",
                            attempt, args.retries, f.url, exc)
                state.mark(f.url, error=str(exc), files=[])
                if attempt < args.retries:
                    time.sleep(args.retry_delay * attempt)
            except Exception as exc:  # noqa: BLE001 - a bug must not stop the run
                log.exception("Unexpected error for %s: %s", f.url, exc)
                state.mark(f.url, error=f"unexpected: {exc}", files=[])
                if attempt < args.retries:
                    time.sleep(args.retry_delay * attempt)
            except KeyboardInterrupt:
                state.mark(f.url, error="interrupted", files=[])
                raise
        else:
            rec = state.records.get(f.url, {})
            state.mark(f.url, status="failed", error=rec.get("error", "unknown"), files=[])
            consecutive_failures += 1
            log.error("FAILED %s after %d attempts", f.url, args.retries)
            if consecutive_failures >= args.cooldown_failures:
                log.warning(
                    "%d failures in a row - cooling down %ds before continuing "
                    "(HTTP 403 storm?). The run does not stop.",
                    consecutive_failures, args.cooldown_seconds,
                )
                time.sleep(args.cooldown_seconds)


def format_size(n: int) -> str:
    size = float(n)
    for unit in ("B", "KiB", "MiB", "GiB", "TiB"):
        if size < 1024 or unit == "TiB":
            return f"{size:.1f} {unit}"
        size /= 1024
    return f"{n} B"


def write_failed_urls(out_dir: Path, state: State) -> Path:
    dest = out_dir / "failed_urls.txt"
    failures = state.failures()
    dest.write_text("\n".join(failures) + ("\n" if failures else ""), encoding="utf-8")
    return dest


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Collect downloads from a Discord channel (gofile + 0807)."
    )
    parser.add_argument("channel_id", help="Discord channel or thread id")
    parser.add_argument("out_dir", help="directory for downloads and state")
    parser.add_argument("--token", help="Discord token (else $DISCORD_TOKEN or prompt)")
    parser.add_argument("--urls-file", type=Path,
                        help="skip the Discord export; read URLs from this file "
                             "(e.g. failed_urls.txt)")
    parser.add_argument("--password", help="password for password-protected gofile links")
    parser.add_argument("--retries", type=int, default=3, help="attempts per URL (default 3)")
    parser.add_argument("--retry-delay", type=int, default=5,
                        help="base seconds between retries, multiplied per attempt (default 5)")
    parser.add_argument("--cooldown-failures", type=int, default=3,
                        help="consecutive failures before a cooldown pause (default 3)")
    parser.add_argument("--cooldown-seconds", type=int, default=60,
                        help="length of the cooldown pause (default 60)")
    parser.add_argument("--only", choices=("gofile", "0807"),
                        help="process only one host class")
    parser.add_argument("--hash", action="store_true",
                        help="record sha256 of each downloaded file")
    parser.add_argument("--extract", action="store_true",
                        help="unzip downloaded archives into per-archive folders")
    parser.add_argument("--fresh", action="store_true",
                        help="re-export the Discord channel even if export.json exists")
    parser.add_argument("--dry-run", action="store_true",
                        help="export/extract/classify only, download nothing")
    args = parser.parse_args()

    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s %(levelname)-7s %(message)s",
        datefmt="%H:%M:%S",
    )

    out_dir = Path(args.out_dir)
    state = State(out_dir / "state.json")

    channel = args.channel_id
    if args.urls_file:
        found = load_urls_file(args.urls_file)
    else:
        token = resolve_token(args.token)
        export_path = dlp_discord.export_channel(
            channel, out_dir, token, force=args.fresh
        )
        found = dlp_discord.extract_urls(export_path)

    if args.only:
        found = [f for f in found if f.host == args.only]

    buckets: dict[str, list[dlp_discord.FoundUrl]] = {"gofile": [], "0807": []}
    for f in found:
        if f.host in buckets:
            buckets[f.host].append(f)
        else:  # classify() already filtered, but be defensive
            state.mark(f.url, status="failed", error=f"unsupported host: {f.url}")

    unsupported = [
        f for f in found if f.host not in buckets
    ]
    for f in unsupported:
        log.warning("Skipping unsupported URL: %s", f.url)

    manifest = out_dir / "manifest.jsonl"
    source_by_url: dict[str, dict] = {
        f.url: {
            "channel": channel,
            "host": f.host,
            "message_id": f.message_id,
            "timestamp": f.timestamp,
            "author": f.author,
            "manifest": str(manifest),
        }
        for f in found
    }

    log.info(
        "Found %d URLs (%d gofile, %d 0807); %d already done from earlier runs",
        len(found),
        len(buckets["gofile"]),
        len(buckets["0807"]),
        len(found) - len(state.pending([f.url for f in found])),
    )

    if args.dry_run:
        for f in found:
            print(f"[{f.host:6}] {f.url}")
        return

    if buckets["gofile"]:
        run_gofile(buckets["gofile"], args, state, source_by_url)
    if buckets["0807"]:
        run_0807(buckets["0807"], args, state, source_by_url)

    failed_file = write_failed_urls(out_dir, state)

    # Summary
    ok = [u for u, r in state.records.items() if r.get("status") == "ok"]
    bad = state.failures()
    total_files = sum(len(r.get("files") or []) for r in state.records.values())
    total_size = sum(
        (f_.get("size") or 0) for r in state.records.values() for f_ in (r.get("files") or [])
    )
    log.info("=" * 60)
    log.info("Done: %d ok (%d files, %s) / %d failed / %d total seen",
             len(ok), total_files, format_size(total_size), len(bad), len(state.records))
    if bad:
        log.info("Failed URLs written to %s - re-run this tool or use "
                 "python3 dlp.py <id> <out> --urls-file %s to retry them.",
                 failed_file, failed_file.name)
        sys.exit(1)


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        log.warning("Interrupted - state.json is saved; re-run to resume.")
        sys.exit(130)