import math
import platform
import subprocess
from typing import Iterable, Tuple

import numpy as np

from src.primesgen import Sieve


class SegmentedEratosthenes(Sieve):
    """
    Utilizes segmentation to handle numbers beyond the scope of the normal erat by reducing the mem used for the sieve,
    but other values have to be tracked between segments AND this data grows with each segment.

    The prime values can be broadcast, if a queue is provided, thus reducing the prime values stores, but the multiples
    will continue to grow (needs some sort put/get that fronts storage on disk (or something))

    Essentially a copy of Kim Walisch's C++ impl (which is way way way faster)
    https://github.com/kimwalisch/primesieve/wiki/Segmented-sieve-of-Eratosthenes

    but the off loading of the data to queues and/or disk is something I want to investigate more
    """

    @staticmethod
    def get_cache_size() -> int:
        os = platform.system()
        if os == 'Darwin':
            # macOS/Darwin
            raw_str = subprocess.run(['sysctl', 'hw.l1dcachesize'], capture_output=True, text=True, check=True).stdout
            return int(raw_str.strip().split(' ')[-1])
        raise NotImplementedError(f'not implemented for your operating system: {os}')

    def __init__(self, segment_max_size: int = -1, primes_queue=None) -> None:
        self._max_sqrt_val = None
        self._max_val = None
        self._seg_max_size = segment_max_size if segment_max_size > 1 else self.get_cache_size()
        self._multiples = self.Multiples()
        self._primes = self.PrimesFactory.get(primes_queue)

    def _create_sieve(self, max_val: int, seg_offset: int, seg_num: int) -> np.ndarray[Tuple[int]]:
        new_sieve = np.ones(max_val + 1, dtype=bool)
        # bye bye evens
        new_sieve[seg_offset % 2::2] = False
        if seg_num == 1:
            new_sieve[1] = False
        # sieve existing multiples
        for p_val, start_val in self._multiples:
            new_sieve[start_val::p_val] = False
            self._multiples.put(p_val, start_val, max_val)
        return new_sieve

    def primes(self, max_val) -> list[int]:
        if max_val < 3:
            return []
        # prep data for all segments based on input
        self._max_val = max_val
        self._max_sqrt_val = int(math.sqrt(self._max_val)) + 1
        num_segs = int((self._max_val + 1) / self._seg_max_size) + 1
        if self._seg_max_size == self._max_val + 1:
            num_segs = 1

        # sieve each segment
        for seg_num in range(1, num_segs + 1):
            self._sieve_segment(seg_num)

        return self._primes.all()

    def _sieve_segment(self, seg_num: int) -> None:
        seg_offset = (seg_num - 1) * self._seg_max_size
        max_val_seg = min(self._max_val - seg_offset, self._seg_max_size - 1)
        max_sqrt_val_seg = min(max(0, self._max_sqrt_val - seg_offset), self._seg_max_size)
        sieve = self._create_sieve(max_val_seg, seg_offset, seg_num)

        # erat (0 --> max_sqrt_val_seg)
        for p_val in (true_index + seg_offset for true_index in range(0, max_sqrt_val_seg) if sieve[true_index]):
            self._primes.append(p_val)
            p_2_val = p_val * p_val
            p_2_val_seg = p_2_val - seg_offset
            self._multiples.put(p_val, p_2_val_seg, max_val_seg)
            sieve[p_2_val_seg::p_val] = False

        # rest (max_sqrt_val_seg -- > end          and        additional offset of max_sqrt_val_seg)
        for p_val in (true_index + max_sqrt_val_seg + seg_offset for true_index in
                      np.where(sieve[max_sqrt_val_seg::])[0]):
            self._primes.append(p_val)
            p_2_val = p_val * p_val
            p_2_val_seg = p_2_val - seg_offset
            self._multiples.put(p_val, p_2_val_seg, max_val_seg)

    class PrimesFactory:
        """
        just a helper to hide the details of queuing vs list-ing primes
        """

        @staticmethod
        def get(in_queue=None):
            return SegmentedEratosthenes.PrimesFactory.PrimesToReturn() if in_queue is None else SegmentedEratosthenes.PrimesFactory.PrimesToQueue(
                in_queue)

        class PrimesToReturn:
            """
            uses a local array
            """

            def __init__(self): self._primes = []

            def append(self, prime: int) -> None: self._primes.append(prime)

            def all(self) -> list[int]: return self._primes

        class PrimesToQueue:
            """
            uses a queue
            """

            def __init__(self, in_queue):
                self._queue = in_queue
                self._queue_count = 0

            def append(self, prime: int) -> None:
                self._queue.put(prime)
                self._queue_count += 1

            def all(self) -> list[int]:
                self._queue.put(None)
                return []

    class Multiples:
        # offload as it grows....which it will keep doing
        # for now a local dict
        def __init__(self) -> None: self.__multiples = {}

        def __iter__(self) -> Iterable[Tuple[int, int]]: return iter(self.__multiples.items())

        def get(self, p_val: int) -> int: return self.__multiples[p_val]

        def put(self, p_val: int, start_val: int, max_val: int) -> None:
            self.__multiples[p_val] = self._get_start_val_in_next_seg(start_val, p_val, max_val)

        @staticmethod
        def _get_start_val_in_next_seg(start_val: int, step: int, max_val: int) -> int:
            max_val_plus_1_seg = max_val + 1
            # change anything here to make life hell
            return start_val + (-max_val_plus_1_seg if start_val > max_val else (math.floor(
                ((max_val - start_val) / step)) + 1) * step - max_val_plus_1_seg)
