#!/usr/bin/env python3
import argparse
import base64
import json
import sys
import time
import urllib.error
import urllib.parse
import urllib.request
from pathlib import Path


API_URL = "https://surepoly.com/api.php"
DEFAULT_OUTPUT_DIR = Path("D:/")
DEFAULT_INTERVAL_SECONDS = 300


class RateLimitError(RuntimeError):
    def __init__(self, message, wait_seconds):
        super().__init__(message)
        self.wait_seconds = max(1, int(wait_seconds or 1))


def call_surepoly_download_api(email):
    url = f"{API_URL}?action=email_download_files&lang=zh"
    body = json.dumps({"email": email}).encode("utf-8")
    request = urllib.request.Request(
        url,
        data=body,
        headers={
            "Accept": "application/json",
            "Content-Type": "application/json",
            "User-Agent": "SurePolyUserDownloader/1.0",
        },
    )
    try:
        with urllib.request.urlopen(request, timeout=60) as response:
            raw = response.read().decode("utf-8", errors="replace")
    except urllib.error.HTTPError as exc:
        raw = exc.read().decode("utf-8", errors="replace")
        try:
            data = json.loads(raw or "{}")
            message = data.get("message") or raw
        except json.JSONDecodeError:
            data = {}
            message = raw
        if exc.code == 429:
            retry_after = exc.headers.get("Retry-After") or data.get("wait_seconds") or 300
            raise RateLimitError(message, retry_after) from exc
        raise RuntimeError(f"SurePoly refused the download request: {message}") from exc

    try:
        data = json.loads(raw or "{}")
    except json.JSONDecodeError as exc:
        raise RuntimeError("SurePoly returned an invalid response. Please try again later.") from exc
    if data.get("success") is False:
        raise RuntimeError(data.get("message") or "SurePoly refused the download request.")
    return data


def save_files(payload, output_dir):
    output_dir.mkdir(parents=True, exist_ok=True)
    saved = []
    for item in payload.get("files") or []:
        name = Path(str(item.get("name") or "")).name
        if name not in {"tt.txt", "targeter.json"}:
            continue
        content = item.get("content_base64") or ""
        data = base64.b64decode(content)
        path = output_dir / name
        path.write_bytes(data)
        saved.append((path, len(data)))
    if len(saved) != 2:
        raise RuntimeError("SurePoly did not return all required files.")
    return saved


def main():
    parser = argparse.ArgumentParser(description="SurePoly signal file downloader")
    parser.add_argument("--email", help="SurePoly account email")
    parser.add_argument("--output", default=str(DEFAULT_OUTPUT_DIR), help="Output folder, default: D:/")
    parser.add_argument("--interval", type=int, default=DEFAULT_INTERVAL_SECONDS, help="Download interval in seconds, default: 300")
    parser.add_argument("--once", action="store_true", help="Download once and exit")
    args = parser.parse_args()

    email = (args.email or input("Enter your SurePoly account email: ")).strip().lower()
    if "@" not in email:
        print("Please enter a valid email address.")
        return 2

    output_dir = Path(args.output).expanduser()
    interval = max(30, int(args.interval or DEFAULT_INTERVAL_SECONDS))

    first_success = False
    while True:
        try:
            payload = call_surepoly_download_api(email)
            access = payload.get("download_access") or {}
            remaining_days = int(float(access.get("remaining_days") or 0))
            confirmed_amount = float(access.get("confirmed_amount") or 0)
            saved = save_files(payload, output_dir)
        except RateLimitError as exc:
            wait = max(interval, exc.wait_seconds + 2)
            print(time.strftime("[%Y-%m-%d %H:%M:%S]"), str(exc))
            print(f"Server rate limit is active. Waiting {wait} seconds before trying again.")
            if args.once:
                return 2
            try:
                time.sleep(wait)
            except KeyboardInterrupt:
                print("\nStopped by user.")
                return 0
            continue
        except Exception as exc:
            print(str(exc))
            return 1 if not first_success else 2

        first_success = True
        print(time.strftime("[%Y-%m-%d %H:%M:%S]"), f"SurePoly download access confirmed for {email}.")
        print(f"Remaining signal download days: {remaining_days}")
        print(f"Confirmed subscription funds: {confirmed_amount:g} USDC")
        for path, size in saved:
            print(f"Downloaded: {path} ({size} bytes)")

        if args.once:
            return 0
        print(f"Next download in {interval} seconds. Press Ctrl+C to stop.")
        try:
            time.sleep(interval)
        except KeyboardInterrupt:
            print("\nStopped by user.")
            return 0


if __name__ == "__main__":
    sys.exit(main())