#!/usr/bin/env -S uv run
# /// script
# requires-python = ">=3.12"
# dependencies = ["httpx"]
# ///
"""
Keep2Share (k2s.cc) downloader — translated from JDownloader's K2SApi + Keep2ShareCc plugins.

Usage:
  ./k2s_download.py <url> [--user EMAIL] [--pass PASSWORD] [--out DIR]

Premium accounts skip captchas and get full speed.
Free/anonymous downloads require solving an image captcha (URL printed to terminal).
"""

import re
import sys
import os
import time
import argparse
import getpass
from pathlib import Path

import httpx

# ── constants ────────────────────────────────────────────────────────────────

API_BASE = "https://k2s.cc/api/v2"
# Alias domains all resolve to k2s.cc internally
SUPPORTED_DOMAINS = {
    "k2s.cc", "keep2share.cc", "keep2.cc",
    "fileboom.me", "fboom.me",
    "tezfiles.com", "publish2.me",
}

FILE_ID_RE = re.compile(
    r"(?i)/(?:f|file|preview)/(?:info/)?([a-z0-9_\-]{13,})"
)
FOLDER_ID_RE = re.compile(r"(?i)/folder(?:/info)?/([a-z0-9]{13,})")

HEADERS = {
    "User-Agent": "JDownloader",
    "Accept-Language": "en-gb, en;q=0.8",
    "Content-Type": "application/json",
}


# ── helpers ──────────────────────────────────────────────────────────────────

def extract_file_id(url: str) -> str | None:
    m = FILE_ID_RE.search(url)
    if not m:
        return None
    fid = m.group(1)
    # fixContentID: lowercase exactly-13-char IDs that aren't "special"
    if len(fid) == 13 and fid.isalpha():
        fid = fid.lower()
    return fid


def is_folder(url: str) -> bool:
    return bool(FOLDER_ID_RE.search(url))


def _post(session: httpx.Client, endpoint: str, payload: dict, auth_token: str | None) -> dict:
    # auth_token goes in the JSON body, not as a header (per K2SApi.handleDownload bytecode)
    if auth_token:
        payload = {**payload, "auth_token": auth_token}
    r = session.post(f"{API_BASE}{endpoint}", json=payload, headers=HEADERS, timeout=30)
    try:
        return r.json()
    except Exception:
        r.raise_for_status()
        raise


# ── API calls (mirroring K2SApi methods) ─────────────────────────────────────

def api_login(session: httpx.Client, username: str, password: str) -> str:
    """POST /login → returns auth_token."""
    data = _post(session, "/login", {"username": username, "password": password}, None)
    token = data.get("auth_token") or data.get("accessToken")
    if not token:
        raise RuntimeError(f"Login failed: {data.get('message') or data}")
    return token


def api_get_file_status(session: httpx.Client, file_id: str, auth_token: str | None) -> dict:
    """POST /getfilestatus → name, size, is_available, access, md5 …"""
    return _post(session, "/getfilestatus", {"id": file_id}, auth_token)


def api_get_files_info(session: httpx.Client, file_ids: list[str], auth_token: str | None) -> dict:
    """POST /getfilesinfo → bulk file metadata."""
    return _post(session, "/getfilesinfo", {"ids": file_ids}, auth_token)


def api_request_captcha(session: httpx.Client, file_id: str) -> dict:
    """POST /requestcaptcha → captcha_url + captcha_challenge."""
    return _post(session, "/requestcaptcha", {"file_id": file_id}, None)


def api_request_recaptcha(session: httpx.Client, file_id: str) -> dict:
    """POST /requestrecaptcha → re_captcha_challenge (site key etc.)."""
    return _post(session, "/requestrecaptcha", {"file_id": file_id}, None)


def api_get_url(
    session: httpx.Client,
    file_id: str,
    auth_token: str | None,
    captcha_challenge: str | None = None,
    captcha_response: str | None = None,
    free_download_key: str | None = None,
) -> dict:
    """POST /geturl → url (direct download link)."""
    payload: dict = {"file_id": file_id}
    if captcha_challenge:
        payload["captcha_challenge"] = captcha_challenge
        payload["captcha_response"] = captcha_response
    if free_download_key:
        payload["free_download_key"] = free_download_key
    return _post(session, "/geturl", payload, auth_token)


# ── 2captcha ─────────────────────────────────────────────────────────────────

class TwoCaptcha:
    """
    Thin wrapper around the 2captcha API v2.
    Mirrors abstractPluginForCaptchaSolverTwoCaptchaAPIV2 from JDownloader:
      POST /createTask  → taskId
      POST /getTaskResult (poll)  → solution.text
    """

    API = "https://api.2captcha.com"
    POLL_INTERVAL = 5   # seconds between getTaskResult polls
    POLL_TIMEOUT  = 120 # seconds before giving up

    def __init__(self, api_key: str, client: httpx.Client):
        self.api_key = api_key
        self.client  = client

    def _call(self, endpoint: str, payload: dict) -> dict:
        r = self.client.post(
            f"{self.API}{endpoint}",
            json={"clientKey": self.api_key, **payload},
            headers={"Content-Type": "application/json"},
            timeout=30,
        )
        r.raise_for_status()
        data = r.json()
        err = data.get("errorId") or data.get("errorCode")
        if err and err != 0:
            raise RuntimeError(f"2captcha error: {data.get('errorCode') or data.get('errorDescription') or data}")
        return data

    def solve_image(self, image_url: str) -> str:
        """Download image from url, submit as ImageToTextTask, return solution text."""
        img_r = self.client.get(image_url, timeout=30)
        img_r.raise_for_status()
        b64 = __import__("base64").b64encode(img_r.content).decode()

        task_id = self._call("/createTask", {"task": {"type": "ImageToTextTask", "body": b64}})["taskId"]
        print(f"  2captcha task {task_id} submitted, polling…", flush=True)

        deadline = time.monotonic() + self.POLL_TIMEOUT
        while time.monotonic() < deadline:
            time.sleep(self.POLL_INTERVAL)
            result = self._call("/getTaskResult", {"taskId": task_id})
            if result.get("status") == "ready":
                text = result["solution"]["text"]
                print(f"  2captcha solved: {text}")
                return text
            print("  …waiting for solution")

        raise RuntimeError(f"2captcha timed out after {self.POLL_TIMEOUT}s")


# ── captcha handling ──────────────────────────────────────────────────────────

def solve_captcha(session: httpx.Client, file_id: str, solver: TwoCaptcha | None) -> tuple[str, str]:
    """
    Request a K2S image captcha and return (challenge, response).
    Uses 2captcha automatically when a solver is provided, otherwise prompts the user.
    """
    cap = api_request_captcha(session, file_id)
    captcha_url = cap.get("captcha_url") or cap.get("url")
    challenge   = cap.get("captcha_challenge") or cap.get("challenge")
    if not captcha_url or not challenge:
        raise RuntimeError(f"Unexpected captcha response: {cap}")

    if solver:
        response = solver.solve_image(captcha_url)
    else:
        print(f"\nCaptcha required. Open this URL in your browser and read the code:")
        print(f"  {captcha_url}\n")
        response = input("Enter captcha text: ").strip()

    return challenge, response


# ── download ──────────────────────────────────────────────────────────────────

def stream_download(session: httpx.Client, url: str, dest: Path, filename: str):
    dest.mkdir(parents=True, exist_ok=True)
    out_path = dest / filename
    print(f"Downloading → {out_path}")
    with session.stream("GET", url, headers=HEADERS, timeout=None) as r:
        r.raise_for_status()
        total = int(r.headers.get("content-length", 0))
        done = 0
        with open(out_path, "wb") as f:
            for chunk in r.iter_bytes(chunk_size=65536):
                f.write(chunk)
                done += len(chunk)
                if total:
                    pct = done * 100 // total
                    print(f"\r  {pct}%  {done // 1024} / {total // 1024} KB", end="", flush=True)
    print(f"\nDone: {out_path}")


# ── main flow ─────────────────────────────────────────────────────────────────

def download(url: str, username: str | None, password: str | None, out_dir: Path, twocaptcha_key: str | None = None):
    if is_folder(url):
        sys.exit("Folder URLs are not supported — add individual file links.")

    file_id = extract_file_id(url)
    if not file_id:
        sys.exit(f"Could not extract file ID from: {url}")

    session = httpx.Client(follow_redirects=True)
    auth_token: str | None = None
    solver = TwoCaptcha(twocaptcha_key, session) if twocaptcha_key else None

    # ── 1. authenticate ───────────────────────────────────────────────────────
    if username and password:
        print(f"Logging in as {username} …")
        auth_token = api_login(session, username, password)
        print("Logged in OK.")

    # ── 2. check file ─────────────────────────────────────────────────────────
    print(f"Checking file {file_id} …")
    info = api_get_file_status(session, file_id, auth_token)

    if not info.get("is_available") and info.get("status") != "success":
        err = info.get("message") or info.get("error") or info
        sys.exit(f"File not available: {err}")

    name = info.get("name") or file_id
    size = info.get("size", 0)
    access = info.get("access", "")
    print(f"  Name   : {name}")
    print(f"  Size   : {size // 1024 // 1024} MB" if size else "  Size   : unknown")
    print(f"  Access : {access}")

    if access == "private" and not auth_token:
        sys.exit("Private file — login required.")

    # ── 3. get download URL ───────────────────────────────────────────────────
    result: dict = {}
    direct_url: str | None = None

    if auth_token:
        # Premium / logged-in path: no captcha needed
        result = api_get_url(session, file_id, auth_token)
        direct_url = result.get("url") or result.get("premlink") or result.get("freelink1")
    else:
        # Free / anonymous path — mirrors K2SApi.handleDownload free flow:
        #   1. solve captcha
        #   2. POST /geturl → may return free_download_key + time_wait instead of url
        #   3. sleep time_wait, re-POST with free_download_key (no captcha) until key is gone
        MAX_CAPTCHA_TRIES = 3
        for attempt in range(MAX_CAPTCHA_TRIES):
            challenge, cap_response = solve_captcha(session, file_id, solver)
            result = api_get_url(session, file_id, None,
                                 captcha_challenge=challenge,
                                 captcha_response=cap_response)
            err_code = result.get("errorCode") or result.get("code")
            if err_code in (30, 33):
                print("Wrong captcha, try again…")
                continue

            # Wait loop: server issues a free_download_key and a time_wait before
            # it hands over the real url (JDownloader allows up to 4 re-polls).
            MAX_WAITS = 4
            for wait_count in range(MAX_WAITS + 1):
                free_key = result.get("free_download_key")
                if not free_key:
                    break  # server ready — url should be present now

                wait_secs = int((result.get("time_wait") or 0))
                if wait_secs > 180:
                    sys.exit(f"Rate limited — server wants {wait_secs}s wait, too long.")
                if wait_count >= MAX_WAITS:
                    sys.exit("Too many wait loops from server, giving up.")

                print(f"  Server asked to wait {wait_secs}s before download is ready…")
                time.sleep(wait_secs)
                result = api_get_url(session, file_id, None, free_download_key=free_key)

            direct_url = result.get("url") or result.get("freelink1")
            if direct_url:
                break
            sys.exit(f"Unexpected response after wait loop: {result}")

    if not direct_url:
        err = result.get("message") or result.get("error") or result
        sys.exit(f"Could not get download URL: {err}")

    # ── 4. download ───────────────────────────────────────────────────────────
    stream_download(session, direct_url, out_dir, name)


# ── CLI ───────────────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(description="Keep2Share downloader")
    parser.add_argument("url", help="k2s.cc / keep2share.cc file URL")
    parser.add_argument("--user", "-u", default=os.environ.get("K2S_USER"), help="Account email")
    parser.add_argument("--pass", "-p", dest="password",
                        default=os.environ.get("K2S_PASS"), help="Account password")
    parser.add_argument("--out", "-o", default=".", help="Output directory (default: .)")
    parser.add_argument("--2captcha", dest="twocaptcha_key",
                        default=os.environ.get("TWOCAPTCHA_API_KEY"),
                        help="2captcha API key (or set TWOCAPTCHA_API_KEY env var)")
    args = parser.parse_args()

    password = args.password
    if args.user and not password:
        password = getpass.getpass(f"Password for {args.user}: ")

    download(
        url=args.url,
        username=args.user,
        password=password,
        out_dir=Path(args.out),
        twocaptcha_key=args.twocaptcha_key,
    )


if __name__ == "__main__":
    main()