#!/usr/bin/python3
# -*- coding: utf-8 -*-

import asyncio
import collections.abc
import concurrent.futures
import io
import itertools
import json
from math import isfinite, inf
import os
import random
import shutil
import sys
import tempfile
import time
import niquests

MIN_SPLIT = 262144
CHUNK_SIZE = 16384
COLOURS = ["\x1b[38;5;16m█"]
COLOURS.extend(f"\x1b[38;5;{i}m█" for i in range(232, 256))
COLOURS.append("\x1b[38;5;15m█")

def generate_session(multiplexed=True):
	return niquests.AsyncSession(
		multiplexed=multiplexed,
		pool_connections=16,
		pool_maxsize=16,
		happy_eyeballs=True,
	)
def header():
	return {
		"Accept": "*/*",
		"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:143.0) Gecko/20100101 Firefox/143.0 AppleWebKit/537.36 Chrome/134.0.0.0 Safari/537.36 Edg/134.0.3124.85",
		"DNT": "1",
		"X-Forwarded-For": "34." + ".".join(str(random.randint(1, 254)) for _ in range(3)),
	}
def box(i):
	if i < 0:
		return "\x1b[38;5;196m█"
	elif not i <= 1:
		return "\x1b[38;5;46m█"
	return COLOURS[round(i * (len(COLOURS) - 1))]
def time_disp(s, rounded=True):
	if not isfinite(s):
		return str(s)
	if rounded:
		s = round(s)
	output = str(s % 60)
	if len(output) < 2:
		output = "0" + output
	if s >= 60:
		temp = str((s // 60) % 60)
		if len(temp) < 2 and s >= 3600:
			temp = "0" + temp
		output = temp + ":" + output
		if s >= 3600:
			temp = str((s // 3600) % 24)
			if len(temp) < 2 and s >= 86400:
				temp = "0" + temp
			output = temp + ":" + output
			if s >= 86400:
				output = str(s // 86400) + ":" + output
	else:
		output = "0:" + output
	return output
def calc_bps(bps):
	for suffix in ("bps", "kbps", "Mbps", "Gbps", "Tbps", "Pbps", "Ebps", "Zbps", "Ybps"):
		bps = round(bps, 4)
		if bps < 1000:
			if bps.is_integer():
				bps = int(bps)
			return f"{bps} {suffix}"
		bps /= 1000
	return "ERR"
def sample(arr, n):
	while len(arr) > n * 2:
		arr = [(a + b) / 2 for a, b in zip(arr[::2], arr[1::2])]
	if len(arr) > n:
		indices = [x * len(arr) / n for x in range(n)]
		arr.append(arr[-1])
		return [arr[int(i)] * (1 - i % 1) + arr[int(i) + 1] * (i % 1) for i in indices]
	return arr


class ChunkManager:

	def __init__(self, url: str, method: str = "get", headers: dict = {}, data: bytes | None = None, filename: str | None = None, fileobj: io.BufferedReader | None = None, concurrent_limit: int = 64, size_limit: int = 1099511627776, verify: bool | None = None, debug: bool = False, log_progress: bool = True, timeout: float | None = None, max_attempts: float = inf):
		self.url = url
		self.method = method
		self.headers = header()
		self.headers.update(headers)
		self.data = data or None
		self.filename = filename
		self.fileobj = fileobj
		self.concurrent_limit = concurrent_limit
		self.size_limit = size_limit
		self.verify = verify
		self.debug = debug
		self.log_progress = log_progress
		self.timeout = timeout
		self.max_attempts = max_attempts
		self.multiplexed = self.concurrent_limit > 1
		self.allow_range_ends = True
		self.timestamp = 0
		self.last_update = 0
		self.latency = inf
		self.workers = []
		self.sessions = set()
		self.request_count = 0
		self.max_bps = 0
		self.max_single_bps = 0
		self.response_headers = {}
		self.status_code = 0

	async def probe_request(self) -> tuple:
		self.session = session = generate_session(self.multiplexed)
		self.sessions.add(session)
		resp = None
		it = None
		try:
			verify = self.verify if self.verify is not None else True
			t = time.perf_counter()
			probe_headers = dict(self.headers)
			probe_headers["Range"] = "bytes=0-"
			req = session.request(
				self.method,
				self.url,
				headers=probe_headers,
				data=self.data,
				stream=True,
				verify=verify,
				timeout=min(3, self.timeout or 3),
			)
			resp = None
			try:
				self.request_count += 1
				resp = await asyncio.wait_for(req, timeout=4)
				ait = resp.iter_content(CHUNK_SIZE)
				if self.timeout:
					ait = asyncio.wait_for(ait, timeout=self.timeout)
				it = await ait
				assert resp.status_code != 416
			except (
				AssertionError,
				TimeoutError,
				asyncio.TimeoutError,
				niquests.ConnectionError,
				niquests.ConnectTimeout,
				niquests.ReadTimeout,
				niquests.Timeout,
				niquests.exceptions.MultiplexingError,
				niquests.exceptions.ChunkedEncodingError,
				AttributeError,
			):
				if resp is not None:
					if it is not None:
						await it.aclose()
					await resp.close()
				self.multiplexed = False
				self.session = session = generate_session(self.multiplexed)
				self.sessions.add(session)
				verify = self.verify or False
				req = session.request(
					self.method,
					self.url,
					headers=self.headers,
					data=self.data,
					stream=True,
					verify=verify,
					timeout=self.timeout,
				)
				if self.timeout:
					req = asyncio.wait_for(req, timeout=self.timeout + 1)
				self.request_count += 1
				resp = await req
				ait = resp.iter_content(CHUNK_SIZE)
				if self.timeout:
					ait = asyncio.wait_for(ait, timeout=self.timeout)
				it = await ait
			self.verify = verify
			self.status_code = resp.status_code
			self.response_headers.update(resp.headers)
			resp.raise_for_status()
			self.latency = time.perf_counter() - t
			return resp, it
		except Exception:
			if resp is not None:
				await resp.close()
			await self.aclose()
			raise

	def update_progress(self, force=False) -> float | None:
		if not self.log_progress:
			return
		ct = time.perf_counter()
		if not force and ct - self.last_update < 0.03:
			return
		self.last_update = ct
		maxbar = 64
		samples = [box(n) for n in sample([2 if worker.done else worker.pos / worker.size for worker in self.workers], maxbar)]
		for i, worker in enumerate(self.workers):
			if worker.error:
				si = int(i / len(self.workers) * len(samples))
				samples[si] = ("\x1b[38;5;226m" if worker.pos > 0 else "\x1b[38;5;196m") + "█"
		s = "\033[K" + "".join(samples)
		dt = max(0.001, ct - self.timestamp)
		timer = time_disp(dt)
		progress = 1 if self.done else self.progress / self.size if self.size > 0 else 0
		percentage = round(progress * 100, 4)
		if percentage.is_integer():
			percentage = int(percentage)
		bpst = calc_bps(self.bps)
		completed = sum(not worker.is_running() for worker in self.workers)
		chunk_count = len(self.workers)
		s2 = f" {timer} {completed}/{chunk_count} ({percentage}%, {bpst})"
		s += "\x1b[38;5;7m" + s2
		print(s, end="\r")
		return progress

	async def join(self, close=False) -> io.BufferedIOBase | str:
		futs = []
		self.workers[0].do(stream=True)
		shattering = asyncio.create_task(self.shatter())
		futs.append(shattering)
		if self.fileobj is not None:
			fp = self.fileobj
			is_stream = True
		else:
			disk = os.path.splitdrive(self.filename or ".")[0]
			folder = disk + "/.temp"
			os.makedirs(folder, exist_ok=True)
			fp = tempfile.NamedTemporaryFile(dir=folder, delete=False)
			is_stream = False
		assert isinstance(fp, (io.BufferedWriter, tempfile._TemporaryFileWrapper))
		with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
			try:
				try:
					assert fp.seekable() and os.path.exists(fp.name)
				except (AssertionError, AttributeError):
					pass
				else:
					fp.seek(0)
				index = 0
				while index < len(self.workers):
					self.workers.sort(key=lambda worker: worker.range.start)
					worker = self.workers[index]
					worker.is_target = True
					stream = index == 0
					fc = worker.do()
					if stream:
						async for b in fc:
							await asyncio.get_event_loop().run_in_executor(executor, fp.write, b)
					else:
						futs.append(fc)
						fc = await fc
						with fc:
							await asyncio.get_event_loop().run_in_executor(executor, shutil.copyfileobj, fc, fp)
					worker.done = True
					index += 1
				if not is_stream:
					fp.close()
					assert self.filename
					try:
						os.replace(fp.name, self.filename)
					except (PermissionError, OSError):
						shutil.copyfile(fp.name, self.filename)
					asize = os.path.exists(self.filename) and os.path.getsize(self.filename)
					assert self.size <= 0 or asize == self.size, f"Expected {self.size} bytes, received {asize}"
			finally:
				while futs:
					await asyncio.wait_for(futs.pop(0), timeout=1)
				if not is_stream and os.path.exists(fp.name):
					fp.close()
					os.remove(fp.name)
				await self.aclose()
				self.update_progress(force=True)
		if close:
			if is_stream:
				fp.close()
				return fp.name
			assert self.filename
			return self.filename
		if is_stream:
			return fp
		assert self.filename
		return open(self.filename, "rb")

	async def shatter(self):
		if not self.concurrent_limit or self.size <= 0:
			return
		worker = self.workers[0]
		while self.workers:
			await asyncio.sleep(0.25)
			remaining = [worker for worker in self.workers if worker.is_running()]
			if not remaining:
				break
			if len(remaining) >= self.concurrent_limit:
				continue
			bps_ratio = self.bps / max(self.max_bps, 1)
			multiplier = ((self.size - self.progress) / self.size)
			try:
				multiplier /= bps_ratio
			except ZeroDivisionError:
				pass
			remaining.sort(key=lambda worker: worker.split_priority(multiplier))
			worker = remaining[-1]
			if worker.split_priority(multiplier) <= 0:
				if bps_ratio < 0.125 and worker.bps / max(1, self.max_single_bps) < 0.5 and time.perf_counter() - worker.timestamp > worker.restart_cooldown:
					worker.needs_restart = True
					worker.restart_cooldown *= 2
				continue
			avg_bps = sum(worker.bps for worker in remaining) / len(remaining)
			ratio = max(1 / 64, min(1 / 2, worker.bps / (avg_bps + worker.bps)))
			worker.split(ratio)

	async def start(self, close=False) -> io.BufferedIOBase | str:
		self.timestamp = time.perf_counter()
		resp, it = await self.probe_request()
		if not self.filename or self.filename.replace("\\", "/").endswith("/"):
			if self.filename:
				path = os.path.abspath(self.filename)
				os.makedirs(path, exist_ok=True)
			else:
				path = os.path.abspath(os.curdir)
			filename: str = (
				resp.headers.get("attachment-filename")
				or resp.headers.get("content-disposition", "").split("filename=", 1)[-1].lstrip('"').split('"', 1)[0].strip().strip('"').strip("'")
				or self.url.rstrip("/").rsplit("/", 1)[-1].split("?", 1)[0]
			)
			import re
			filename = re.sub(r'[<>:"/\\|?*\x00-\x1F]', "_", filename)
			if "." not in filename:
				ctype = resp.headers.get("content-type")
				if ctype and ctype != "application/octet-stream":
					import mimetypes
					ext = mimetypes.guess_extension(ctype)
					if ext:
						filename = filename + "." + ext
			self.filename = os.path.join(path, filename)
		elif not self.fileobj and self.filename == "-":
			self.fileobj = (sys.__stdout__ or sys.stdout).buffer
			sys.stdout = sys.stderr
		try:
			self.size = int(resp.headers.get("content-length") or resp.headers["content-range"].rsplit("/", 1)[-1])
		except (KeyError, ValueError):
			self.size = 0
		if (
			self.size <= 0
			or "content-range" not in resp.headers and "bytes" not in resp.headers.get("accept-ranges", "").lower()
			or resp.headers.get("content-encoding") not in (None, "identity")
		):
			self.concurrent_limit = 0
		worker = ChunkWorker(
			ctx=self,
			resp=resp,
			it=it,
			url=resp.url,
		)
		self.workers.append(worker)
		return await self.join(close=close)

	def run(self, close=False, return_headers=False):
		# import selectors
		# loop = asyncio.SelectorEventLoop(selectors.SelectSelector())
		loop = asyncio.new_event_loop()
		resp = loop.run_until_complete(self.start(close=close))
		if return_headers:
			return resp, self.response_headers
		return resp

	async def aclose(self):
		for session in self.sessions:
			await session.close()

	@property
	def done(self):
		return not self.workers or self.workers[-1].done
	@property
	def progress(self):
		return sum(worker.pos for worker in self.workers)
	@property
	def bps(self):
		if self.done:
			bps = self.size * 8 / max(0.001, time.perf_counter() - self.timestamp)
		else:
			bpss = [worker.bps for worker in self.workers]
			bps = sum(bpss)
			self.max_single_bps = max(self.max_single_bps, max(bpss))
		self.max_bps = max(self.max_bps, bps)
		return bps


_range = range

class ChunkWorker:

	def __init__(self, ctx: ChunkManager, range: _range | None = None, resp: niquests.AsyncResponse | None = None, it: collections.abc.AsyncGenerator | None = None, url: str | None = None):
		self.ctx = ctx
		self.running = None
		headers = ctx.headers
		if range is None:
			n = min(ctx.size or inf, ctx.size_limit)
			assert isinstance(n, int)
			range = _range(0, n)
		self.range = range
		self.resp = resp
		self.it = it
		self.url = url
		self.headers = headers
		self.pos = 0
		self.timestamp = time.perf_counter()
		self.latency = ctx.latency
		self.last_split = self.timestamp
		self.ttfr = inf
		self.recv_times = []
		self.done = False
		self.error = None
		self.needs_restart = False
		self.restart_cooldown = 5
		self.is_target = False

	async def refresh_resp(self, attempt=0, timeout=None) -> niquests.AsyncResponse:
		if self.resp is None:
			ctx = self.ctx
			self.timestamp = time.perf_counter()
			self.last_split = self.timestamp
			self.ttfr = inf
			if attempt > 1 or not ctx.request_count & 15:
				self.url = ctx.url
				ctx.session = session = generate_session(False if attempt > 2 else ctx.multiplexed)
				ctx.sessions.add(session)
			else:
				session = ctx.session
			headers = self.headers.copy()
			headers["Priority"] = "i"
			if ctx.allow_range_ends:
				headers["Range"] = f"bytes={self.start + self.pos}-{self.end - 1}"
			else:
				headers["Range"] = f"bytes={self.start + self.pos}-"
			req = session.request(
				ctx.method,
				self.url,
				headers=headers,
				data=ctx.data,
				stream=True,
				verify=ctx.verify,
				timeout=timeout,
			)
			if timeout:
				req = asyncio.wait_for(req, timeout=timeout + 1)
			ctx.request_count += 1
			self.resp = await req
			ait = self.resp.iter_content(CHUNK_SIZE if attempt else CHUNK_SIZE * 8)
			if timeout:
				ait = asyncio.wait_for(ait, timeout=timeout)
			self.it = await ait
			if "Range" in self.headers:
				range_sent = self.headers["Range"].split("=", 1)[-1].split("-", 1)[0]
				range_recv = self.resp.headers["content-range"].split("/", 1)[0].split(None, 1)[-1].split("-", 1)[0]
				assert range_recv == range_sent, "Server failed to serve range header as specified!"
			self.latency = time.perf_counter() - self.timestamp
			self.error = None
			self.needs_restart = False
		self.resp.raise_for_status()
		return self.resp

	def split(self, ratio=0.5):
		offset = self.pos + round(self.remainder * ratio)
		cur_range = range(self.start, self.start + offset)
		new_range = range(self.start + offset, self.end)
		if len(cur_range) < MIN_SPLIT or len(new_range) < MIN_SPLIT:
			self.last_split = inf
			return
		worker = ChunkWorker(
			ctx=self.ctx,
			url=self.url,
			range=new_range,
		)
		self.last_split = worker.last_split
		self.range = cur_range
		self.ctx.workers.append(worker)
		worker.do()

	async def stream(self):
		ctx = self.ctx
		for attempt in itertools.count(1):
			timeout = attempt * 3 + self.latency
			delay = attempt ** 2 + 1 + random.random()
			t = time.perf_counter()
			try:
				await self.refresh_resp(attempt=attempt, timeout=timeout)
				assert self.it is not None
				try:
					while True:
						t = time.perf_counter()
						fut = self.it.__anext__()
						content = await asyncio.wait_for(fut, timeout=timeout)
						expected = self.remainder
						if len(content) > expected:
							content = memoryview(content)[:expected]
						length = len(content)
						if length <= 0:
							break
						yield content
						self.pos += length
						self.recv_times.append((t, length))
						if self.pos >= self.size:
							break
						if not isfinite(self.ttfr):
							self.ttfr = time.perf_counter() - self.timestamp
						ctx.update_progress()
						if self.needs_restart:
							raise AttributeError
				except (StopIteration, StopAsyncIteration):
					pass
				if self.remainder <= 0:
					break
			except (
				TimeoutError,
				asyncio.TimeoutError,
				niquests.ConnectionError,
				niquests.ConnectTimeout,
				niquests.ReadTimeout,
				niquests.Timeout,
				niquests.exceptions.ChunkedEncodingError,
				AttributeError,
			) as ex:
				if ctx.debug:
					print(repr(ex))
				self.error = ex
			except niquests.HTTPError as ex:
				assert ex.response
				if ex.response.status_code == 416:
					ctx.allow_range_ends = False
				elif ex.response.status_code not in (408, 420, 425, 429, 460, 502, 503, 504, 508, 509, 522, 524, 529, 599):
					raise
				if ctx.debug:
					print(repr(ex))
				self.error = ex
				await asyncio.sleep(2 ** attempt)
			finally:
				if self.resp is not None:
					if self.it is not None:
						await self.it.aclose()
					await self.resp.close()
					self.resp = None
				ctx.update_progress(force=True)
			if attempt + 1 >= ctx.max_attempts:
				raise TimeoutError("Maximum attempts exceeded for worker.")
			sleep = delay + t - time.perf_counter()
			if sleep > 0:
				await asyncio.sleep(sleep)

	async def run(self) -> tempfile.SpooledTemporaryFile:
		self.fp = tempfile.SpooledTemporaryFile(max_size=12 * 1048576)
		async for b in self.stream():
			self.fp.write(b)
		self.fp.flush()
		self.fp.seek(0)
		return self.fp

	def do(self, stream=False) -> collections.abc.Awaitable:
		if self.running:
			return self.running
		if stream:
			fut = self.stream()
			self.running = fut
		else:
			fut = self.run()
			self.running = asyncio.create_task(fut)
		return self.running

	def split_priority(self, multiplier=1) -> float:
		if self.error or self.size <= MIN_SPLIT:
			return 0
		ctx = self.ctx
		ct = time.perf_counter()
		if ct - self.last_split < self.ttfr * 2.5:
			return 0
		concurrency = sum((not worker.done) + bool(worker.error) for worker in ctx.workers)
		if (ct - self.last_split - self.ttfr) * multiplier < concurrency ** 2 / 4:
			return 0
		bps = self.bps
		if bps <= 0:
			return 0
		if self.remainder / self.bps * 8 <= self.ttfr * 2 + 1:
			return 0
		return (ct - self.last_split - self.ttfr) / self.bps * self.remainder / self.size * (1 + self.is_target)

	@property
	def start(self):
		return self.range.start
	@property
	def end(self):
		return self.range.stop
	@property
	def size(self):
		return self.range.stop - self.range.start
	@property
	def remainder(self):
		return self.range.stop - self.range.start - self.pos
	@property
	def bps(self):
		if len(self.recv_times) >= 128 and not len(self.recv_times) & 1:
			recv_times = []
			for i in range(0, len(self.recv_times), 2):
				t1, n1 = self.recv_times[i]
				t2, n2 = self.recv_times[i + 1]
				recv_times.append((t1, n1 + n2))
			self.recv_times = recv_times
		t = time.perf_counter()
		while self.recv_times and t > self.recv_times[0][0] + 5:
			self.recv_times.pop(0)
		if not self.recv_times:
			return 0
		return sum(t[-1] for t in self.recv_times) * 8 / max(0.001, t - self.recv_times[0][0])
	def is_running(self):
		if not self.running:
			return False
		if hasattr(self.running, "done"):
			return not self.running.done()
		return not self.done


async def shatter_request(url, method="get", headers={}, data=None, filename=None, fileobj=None, concurrent_limit=64, size_limit=1099511627776, verify=None, debug=False, log_progress=True, timeout=None, max_attempts=inf, return_headers=False):
	ctx = ChunkManager(
		url=url,
		method=method,
		headers=headers,
		data=data,
		filename=filename,
		fileobj=fileobj,
		concurrent_limit=concurrent_limit,
		size_limit=size_limit,
		verify=verify,
		debug=debug,
		log_progress=log_progress,
		timeout=timeout,
		max_attempts=max_attempts,
	)
	resp = await ctx.start()
	if return_headers:
		return resp, ctx.response_headers
	return resp
parallel_request = shatter_request

def download(url, method="get", headers={}, data=None, filename=None, fileobj=None, concurrent_limit=64, size_limit=1099511627776, verify=None, debug=False, log_progress=True, timeout=None, max_attempts=inf, return_headers=False):
	ctx = ChunkManager(
		url=url,
		method=method,
		headers=headers,
		data=data,
		filename=filename,
		fileobj=fileobj,
		concurrent_limit=concurrent_limit,
		size_limit=size_limit,
		verify=verify,
		debug=debug,
		log_progress=log_progress,
		timeout=timeout,
		max_attempts=max_attempts,
	)
	resp = ctx.run()
	if return_headers:
		return resp, ctx.response_headers
	return resp


try:
	from importlib.metadata import version
	__version__ = version("streamshatter")
except Exception:
	__version__ = "0.0.0-unknown"

def main():
	import argparse
	parser = argparse.ArgumentParser(
		prog="streamshatter",
		description="Multiplexed chunked file downloader",
	)
	parser.add_argument("-V", '--version', action='version', version=f'%(prog)s {__version__}')
	parser.add_argument("-H", '--headers', help="HTTP headers, interpreted as JSON", required=False, default="{}")
	parser.add_argument("-l", "-cl", '--concurrent-limit', help="Limits the amount of concurrent requests; defaults to 64", type=int, required=False, default=64)
	parser.add_argument("-sl", '--size-limit', help="Limits the amount of data to download; defaults to 1099511627776", type=int, required=False, default=1099511627776)
	parser.add_argument("-t", '--timeout', help="Limits the amount of time allowed for the initial request to succeed", type=float, required=False, default=60)
	parser.add_argument("-ma", '--max-attempts', help="Limits the amount of retries allowed for each request", type=float, required=False, default=inf)
	parser.add_argument("-s", "--ssl", action=argparse.BooleanOptionalAction, default=True, help="Enforces SSL verification; defaults to TRUE")
	parser.add_argument("-d", "--debug", action=argparse.BooleanOptionalAction, default=False, help="Terminates immediately upon non-timeout errors, and writes the response data for errored chunks; defaults to FALSE")
	parser.add_argument("-lp", "--log-progress", action=argparse.BooleanOptionalAction, default=True, help="Continually updates a progress bar in the standard output; defaults to TRUE")
	parser.add_argument("url", help="Target URL", nargs="?", default="")
	parser.add_argument("filename", help='Output filename; use "-" for stdout pipe', nargs="?", default="")
	args = parser.parse_args()
	if not args.url:
		args.url = input("Please use `streamshatter -h` for help, or input a URL to download directly: ").strip()
	if os.name == "nt":
		os.system("color")
	download(
		url=args.url,
		headers=json.loads(args.headers),
		filename=args.filename,
		concurrent_limit=args.concurrent_limit,
		size_limit=args.size_limit,
		verify=args.ssl,
		debug=args.debug,
		log_progress=args.log_progress,
		timeout=args.timeout,
		max_attempts=args.max_attempts,
	)

if __name__ == "__main__":
	main()