#!/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.

Already-downloaded URLs are recorded in a global store (default
~/.local/share/dlphelper/urls.json), so a link that was downloaded for one
channel is never re-downloaded for another - and re-runs of the same channel
skip everything finished (see also --assume-done for importing old lists).

Usage:
    python3 dlp.py <channel_id> <out_dir> [options]
    python3 dlp.py --subs subscriptions.json [options]
    python3 dlp.py <channel_id> <out_dir> --urls-file failed_urls.txt
    python3 dlp.py --assume-done old_urls.txt          # import past downloads
"""

from __future__ import annotations

import argparse
import getpass
import hashlib
import json
import logging
import os
import random
import sys
import time
import zipfile
from itertools import chain
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.globalstore import GlobalStore, default_global_path
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 read_url_lines(path: Path) -> list[str]:
    return [
        line.strip()
        for line in path.read_text(encoding="utf-8").splitlines()
        if line.strip() and not line.strip().startswith("#")
    ]


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 url in read_url_lines(path):
        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,
    gstore: GlobalStore | None,
    url: str,
    files: list[dict],
    source: dict,
) -> None:
    state.mark(url, status="ok", error=None, files=files)
    if gstore is not None:
        gstore.mark(url, channel=source.get("channel"), host=source.get("host"),
                    file_count=len(files))
    for f in files:
        append_manifest(Path(source["manifest"]), {
            "channel": source["channel"],
            "url": url,
            "host": source.get("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],
    is_done,
    gstore: GlobalStore | None,
) -> None:
    if not [f for f in urls if not is_done(f.url)]:
        return  # nothing pending: no guest account creation, no API calls
    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 is_done(f.url):
                continue
            state.mark(f.url, status="failed", error=str(exc), files=[])
        return

    consecutive_failures = 0
    for f in urls:
        if is_done(f.url):
            log.info("Skip (already downloaded): %s", 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, gstore, 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],
    is_done,
    gstore: GlobalStore | None,
) -> None:
    if not [f for f in urls if not is_done(f.url)]:
        return
    out_dir = Path(args.out_dir) / "downloads" / "0807"
    consecutive_failures = 0
    for f in urls:
        if is_done(f.url):
            log.info("Skip (already downloaded): %s", 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, gstore, 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"
    dest.parent.mkdir(parents=True, exist_ok=True)
    failures = state.failures()
    dest.write_text("\n".join(failures) + ("\n" if failures else ""), encoding="utf-8")
    return dest


def mark_assume_done(path: Path, gstore: GlobalStore | None) -> int:
    """Record URLs as already downloaded (no files) so they are never fetched.

    Manual import step for the pre-dlp.py era: point it at an old URL list you
    know is fully downloaded. There is no verification - garbage in, skipped
    forever.
    """
    if gstore is None:
        raise SystemExit("--assume-done needs the global store (drop --no-global)")
    count = 0
    for url in read_url_lines(path):
        gstore.mark(url, channel=None, host=dlp_discord.classify(url), file_count=0,
                    note="assumed done via --assume-done")
        count += 1
    log.info("Marked %d URLs as done in the global store (%s)", count, gstore.path)
    return count


def process_channel(
    channel_id: str,
    out_dir: Path,
    args: argparse.Namespace,
    token: str | None,
    gstore: GlobalStore | None,
    force_export: bool = False,
) -> int:
    """Run the full pipeline for one channel; returns the number of failures."""
    state = State(out_dir / "state.json")
    manifest = out_dir / "manifest.jsonl"

    def is_done(url: str) -> bool:
        return state.is_done(url) or (gstore is not None and gstore.is_done(url))

    if args.urls_file:
        found = load_urls_file(args.urls_file)
    else:
        if token is None:
            raise SystemExit("internal error: no token resolved")
        export_path = dlp_discord.export_channel(
            channel_id, out_dir, token, force=force_export or 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
            log.warning("Skipping unsupported URL: %s", f.url)
            state.mark(f.url, status="failed", error="unsupported host")

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

    # --limit applies per host class to the *pending* URLs, so a test run
    # downloads exactly N new things per host (done ones don't use up the quota)
    if args.limit is not None:
        for key, items in buckets.items():
            pending = [f for f in items if not is_done(f.url)]
            if args.shuffle:
                random.shuffle(pending)
            buckets[key] = pending[: args.limit]

    pending = [f for f in found if not is_done(f.url)]
    if args.limit is not None:
        pending = [f for f in chain.from_iterable(buckets.values()) if not is_done(f.url)]
    log.info(
        "%s: %d URLs (%d gofile, %d 0807); %d done, %d to download",
        channel_id, len(found),
        len(buckets["gofile"]), len(buckets["0807"]),
        len(found) - len(pending), len(pending),
    )

    if args.dry_run:
        selected = {f.url for f in chain.from_iterable(buckets.values())}
        for f in found:
            if f.url in selected:
                flag = "download"
            elif is_done(f.url):
                flag = "done"
            else:
                flag = "skip"  # cut by --limit / --only
            print(f"[{f.host:6}] [{flag:8}] {f.url}")
        return 0

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

    failed_file = write_failed_urls(out_dir, state)

    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("%s: %d ok (%d files, %s) / %d failed / %d seen",
             out_dir, 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)
    return len(bad)


def load_subscriptions(path: Path) -> list[dict]:
    """Load subscriptions.json: [{"name": ..., "channel_id": ..., "out_dir": ...}]."""
    data = json.loads(path.read_text(encoding="utf-8"))
    entries = data.get("subscriptions", []) if isinstance(data, dict) else data
    if not isinstance(entries, list):
        raise SystemExit(f"{path}: expected a list or {'{'}\"subscriptions\": [...]{'}'}")
    subs = []
    for entry in entries:
        if not isinstance(entry, dict) or "channel_id" not in entry or "out_dir" not in entry:
            raise SystemExit(
                f"{path}: each subscription needs 'channel_id' and 'out_dir' "
                f"(got: {entry})"
            )
        subs.append(entry)
    return subs


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Collect downloads from Discord channels (gofile + 0807)."
    )
    parser.add_argument("channel_id", nargs="?", help="Discord channel or thread id")
    parser.add_argument("out_dir", nargs="?", help="directory for downloads and state")
    parser.add_argument("--subs", type=Path, metavar="FILE",
                        help="process a subscription list instead of one channel "
                             "(entries: channel_id, out_dir, optional name)")
    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("--limit", type=int, metavar="N",
                        help="download at most N pending URLs per host class "
                             "(test runs)")
    parser.add_argument("--shuffle", action="store_true",
                        help="randomize which URLs a --limit picks")
    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 "
                             "(subs mode does this by default)")
    parser.add_argument("--no-fresh", action="store_true",
                        help="subs mode: reuse cached exports instead of re-exporting")
    parser.add_argument("--dry-run", action="store_true",
                        help="export/extract/classify only, download nothing")
    parser.add_argument("--global-state", type=Path, metavar="FILE",
                        help=f"cross-run URL store (default: {default_global_path()})")
    parser.add_argument("--no-global", action="store_true",
                        help="disable the cross-run URL store for this run")
    parser.add_argument("--assume-done", type=Path, metavar="FILE",
                        help="mark every URL in the file as already downloaded, "
                             "then exit (manual import of old download lists)")
    args = parser.parse_args()

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

    gstore = None
    if not args.no_global:
        gstore = GlobalStore(args.global_state or default_global_path())

    if args.assume_done:
        if args.channel_id or args.subs:
            parser.error("--assume-done cannot be combined with a channel or --subs")
        mark_assume_done(args.assume_done, gstore)
        return

    if args.subs:
        if args.channel_id or args.out_dir or args.urls_file:
            parser.error("--subs cannot be combined with channel_id/out_dir/--urls-file")
        subs = load_subscriptions(args.subs)
        token = resolve_token(args.token)
        total_failed = 0
        for entry in subs:
            out_dir = Path(entry["out_dir"])
            log.info("=== subscription: %s ===", entry.get("name") or entry["channel_id"])
            total_failed += process_channel(
                entry["channel_id"], out_dir, args, token, gstore,
                force_export=not args.no_fresh,  # subs mode re-exports by default
            )
        if total_failed:
            log.warning("%d URL(s) failed in total across subscriptions.", total_failed)
            sys.exit(1)
        return

    if not args.channel_id or not args.out_dir:
        parser.error("channel_id and out_dir are required (or use --subs/--assume-done)")

    # A token is only needed when we will export a channel (--urls-file skips that)
    token = None if args.urls_file else resolve_token(args.token)
    failures = process_channel(args.channel_id, Path(args.out_dir), args, token, gstore)
    if failures:
        sys.exit(1)


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