# Copyright (c) 2026 동일비전(Dongil Vision Korea). All Rights Reserved.
"""
╔══════════════════════════════════════════════════════════════════════════════╗
║  ALASKA v2.3 - SmBlock                                                       ║
║  Shared Memory Block Allocator                                               ║
╠══════════════════════════════════════════════════════════════════════════════╣
║  Version : 2.3.1                                                             ║
║  Date    : 2026-02-28                                                        ║
╠══════════════════════════════════════════════════════════════════════════════╣
║  Summary:                                                                    ║
║    프로세스 간 공유 메모리 이미지 블록 풀 관리자                             ║
║    - Free List 알고리즘 O(1) 블록 할당/해제                                  ║
║    - Layout: [Index: int16[N]] + [Data: uint8[N][H][W][C]] + [PID] + [Heads] ║
║    - SmBlockBuffer 컨텍스트 매니저로 자동 해제                               ║
╠══════════════════════════════════════════════════════════════════════════════╣
║  Changes (v2.3.1):                                                           ║
║    - used_count(int32) 필드 추가로 count() O(1) 최적화                       ║
║    - malloc 시 데이터 복사를 lock 외부에서 수행하여 동시성 향상              ║
╚══════════════════════════════════════════════════════════════════════════════╝
"""

import os
from multiprocessing.shared_memory import SharedMemory
from typing import Tuple
import numpy as np


class SmBlock:
    """프로세스 간 공유 메모리 이미지 블록 풀 관리자 (v2.3.1: O(1) count 지원)"""

    __slots__ = ('_name', '_shape', '_maxsize', '_item_size', '_index_size',
                 '_is_new', '_shm', '_index_buffer', '_data_buffer', '_free_head',
                 '_lock', '_last_alloc', '_pid_buffer', '_used_count_buffer')

    def __init__(self, name: str, shape: Tuple[int, ...], maxsize: int,
                 create: bool = False, lock=None):
        """SmBlock 초기화.

        Args:
            name: 공유 메모리 이름 (시스템 전역 고유)
            shape: 이미지 형태 (height, width, channels)
            maxsize: 최대 블록 수 (풀 크기)
            create: True=새 공유 메모리 생성, False=기존에 연결
            lock: multiprocessing.Lock 동시성 제어용

        Raises:
            FileExistsError: create=True인데 이미 존재할 때
            FileNotFoundError: create=False인데 존재하지 않을 때
        """
        self._name = name
        self._shape = shape
        self._maxsize = maxsize
        self._lock = lock
        self._is_new = create
        self._last_alloc = -1

        # 크기 계산
        self._item_size = int(np.prod(shape))  # H * W * C
        self._index_size = 2 * maxsize  # sizeof(int16) * N
        pid_size = 4 * maxsize # sizeof(int32) * N

        # 총 공유 메모리 크기 (Index + Data + PID + FreeHead + UsedCount)
        # v2.3.1: UsedCount(int32) 추가로 count() O(1) 지원
        total_size = self._index_size + (self._item_size * maxsize) + pid_size + 2 + 4

        # 공유 메모리 생성 또는 연결
        if create:
            self._shm = SharedMemory(name=name, create=True, size=total_size)
            self._setup_buffers()
            self._initialize_free_list()
        else:
            self._shm = SharedMemory(name=name, create=False)
            self._setup_buffers()

    def _setup_buffers(self):
        """numpy 배열로 공유 메모리 영역 매핑."""
        buf = self._shm.buf
        n = self._maxsize

        # Index Buffer: offset 0, size 2*N
        self._index_buffer = np.ndarray(
            (n,), dtype=np.int16,
            buffer=buf[0:self._index_size]
        )

        # Data Buffer: offset 2*N, size N*H*W*C
        data_offset = self._index_size
        data_size = self._item_size * n
        self._data_buffer = np.ndarray(
            (n,) + self._shape, dtype=np.uint8,
            buffer=buf[data_offset:data_offset + data_size]
        )

        # PID Buffer: offset 2*N + DataSize, size 4*N (v2.3)
        pid_offset = data_offset + data_size
        self._pid_buffer = np.ndarray(
            (n,), dtype=np.int32,
            buffer=buf[pid_offset:pid_offset + (4 * n)]
        )

        # Free Head: offset 2*N + DataSize + PIDSize, size 2
        free_head_offset = pid_offset + (4 * n)
        self._free_head = np.ndarray(
            (1,), dtype=np.int16,
            buffer=buf[free_head_offset:free_head_offset + 2]
        )

        # Used Count: offset 2*N + DataSize + PIDSize + 2, size 4 (v2.3.1)
        used_count_offset = free_head_offset + 2
        self._used_count_buffer = np.ndarray(
            (1,), dtype=np.int32,
            buffer=buf[used_count_offset:used_count_offset + 4]
        )

    def _initialize_free_list(self):
        """Free List 초기화 (생성 시에만 호출)."""
        n = self._maxsize
        # 각 블록이 다음 블록을 가리키도록 설정
        for i in range(n - 1):
            self._index_buffer[i] = i + 1
        self._index_buffer[n - 1] = -1  # 마지막 블록은 끝 표시
        self._free_head[0] = 0  # 첫 번째 free 블록
        self._used_count_buffer[0] = 0

    def malloc(self, image: np.ndarray) -> int:
        """이미지를 풀에 할당하고 인덱스 반환. O(1) 시간복잡도.

        Args:
            image: 저장할 이미지 (dtype=uint8, shape=생성자의 shape과 동일)

        Returns:
            할당된 블록 인덱스 (0 ~ maxsize-1), 풀이 가득 차면 -1

        Raises:
            ValueError: 이미지 shape이 맞지 않을 때
        """
        if image.shape != self._shape:
            raise ValueError(f"Shape mismatch: expected {self._shape}, got {image.shape}")

        # 1. 인덱스 확보 (Lock 범위 최소화)
        with self._lock:
            head = self._free_head[0]
            if head == -1:
                return -1

            next_free = self._index_buffer[head]
            self._free_head[0] = next_free
            self._index_buffer[head] = -1
            self._used_count_buffer[0] += 1
            idx = int(head)

        # 2. 데이터 복사 및 소유권 기록 (Lock 외부에서 수행하여 동시성 향상)
        # 단, 다른 프로세스가 해제하기 전에 복사가 완료되어야 하므로 
        # 실제로는 이 블록을 소유한 프로세스가 제어권을 가지는 동안 안전함.
        self._data_buffer[idx][:] = image
        self._pid_buffer[idx] = os.getpid()
        self._last_alloc = idx
        return idx

    def malloc2(self, image: np.ndarray) -> int:
        """마지막 할당 위치 다음부터 순차 탐색하여 빈 공간 할당."""
        if image.shape != self._shape:
            raise ValueError(f"Shape mismatch: expected {self._shape}, got {image.shape}")

        with self._lock:
            n = self._maxsize
            start = (self._last_alloc + 1) % n

            for offset in range(n):
                idx = (start + offset) % n
                if self._index_buffer[idx] != -1 or idx == self._free_head[0]:
                    if idx == self._free_head[0]:
                        self._free_head[0] = self._index_buffer[idx]
                    else:
                        prev = self._free_head[0]
                        while prev != -1 and self._index_buffer[prev] != idx:
                            prev = self._index_buffer[prev]
                        if prev != -1:
                            self._index_buffer[prev] = self._index_buffer[idx]

                    self._index_buffer[idx] = -1
                    self._used_count_buffer[0] += 1
                    self._data_buffer[idx][:] = image
                    self._pid_buffer[idx] = os.getpid()
                    self._last_alloc = idx
                    return idx
            return -1

    def alloc(self) -> int:
        """인덱스만 할당 (데이터 복사 없음). O(1) 시간복잡도."""
        with self._lock:
            head = self._free_head[0]
            if head == -1:
                return -1

            next_free = self._index_buffer[head]
            self._free_head[0] = next_free
            self._index_buffer[head] = -1
            self._used_count_buffer[0] += 1
            self._last_alloc = head
            return int(head)

    def get_buffer(self, index: int) -> np.ndarray:
        """할당된 블록의 쓰기 가능한 버퍼 참조 반환 (복사 없음)."""
        if index < 0 or index >= self._maxsize:
            raise IndexError(f"Index {index} out of range")
        return self._data_buffer[index]

    def mfree(self, index: int):
        """할당된 블록을 해제하고 Free List에 반환."""
        if index < 0 or index >= self._maxsize:
            raise IndexError(f"Index {index} out of range")

        with self._lock:
            if self._index_buffer[index] != -1:
                raise ValueError(f"Block {index} is already free")

            self._pid_buffer[index] = 0
            self._index_buffer[index] = self._free_head[0]
            self._free_head[0] = index
            self._used_count_buffer[0] -= 1

    def get(self, index: int) -> np.ndarray:
        """할당된 블록의 이미지 데이터 반환."""
        if index < 0 or index >= self._maxsize:
            raise IndexError(f"Index {index} out of range")

        with self._lock:
            if self._index_buffer[index] != -1:
                raise ValueError(f"Block {index} is free (not allocated)")

            return self._data_buffer[index].copy()

    def size(self) -> int:
        """풀의 최대 크기(블록 수) 반환."""
        return self._maxsize

    def count(self) -> int:
        """현재 사용 가능한(비어있는) 블록 수 반환. O(1) 최적화."""
        return self._maxsize - int(self._used_count_buffer[0])

    def close(self):
        """공유 메모리 연결 해제 및 정리."""
        self._shm.close()
        if self._is_new:
            try:
                self._shm.unlink()
            except FileNotFoundError:
                pass
            except PermissionError:
                import warnings; warnings.warn(f"SHM unlink failed (in use): {self._shm.name}", ResourceWarning, stacklevel=2)

    @property
    def name(self) -> str:
        return self._name

    @property
    def shape(self) -> Tuple[int, ...]:
        return self._shape

    @property
    def maxsize(self) -> int:
        return self._maxsize

    @property
    def item_size(self) -> int:
        return self._item_size

    @property
    def index_size(self) -> int:
        return self._index_size

    @property
    def is_new(self) -> bool:
        return self._is_new

    def buffer(self, copy: bool = False) -> 'SmBlockBuffer':
        """버퍼 컨텍스트 매니저 (with문으로 자동 할당/해제)"""
        return SmBlockBuffer(self, copy=copy)

    def __str__(self) -> str:
        return f"SmBlock[ {self._name}, {self._shape}, {self._maxsize} ]"

    def __repr__(self) -> str:
        return self.__str__()


class SmBlockBuffer:
    """SmBlock 버퍼 컨텍스트 매니저 (with문으로 자동 할당/해제)"""
    __slots__ = ('_smblock', '_index', '_data', '_copy', '_released')

    def __init__(self, smblock: 'SmBlock', copy: bool = False):
        self._smblock = smblock
        self._copy = copy
        self._index = -1
        self._data = None
        self._released = False

    def __enter__(self) -> 'SmBlockBuffer':
        self._index = self._smblock.alloc()
        if self._index < 0:
            raise RuntimeError(f"SmBlock '{self._smblock.name}' is full (no free blocks)")
        buf = self._smblock.get_buffer(self._index)
        self._data = buf.copy() if self._copy else buf
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if not self._released and self._index >= 0:
            try:
                self._smblock.mfree(self._index)
            except ValueError:
                pass  # 이미 해제됨
            self._released = True
        return False

    @property
    def data(self) -> np.ndarray:
        return self._data

    @property
    def index(self) -> int:
        return self._index

    def release(self):
        """수동 해제"""
        if not self._released and self._index >= 0:
            self._smblock.mfree(self._index)
            self._released = True

    def keep(self):
        """자동 해제 방지 (나중에 수동 mfree 필요)"""
        self._released = True

    def __repr__(self) -> str:
        return f"SmBlockBuffer({self._index}, {'X' if self._released else 'O'})"
