"""
╔══════════════════════════════════════════════════════════════════════════════╗
║  ALASKA v2.0 - SmBlock                                                       ║
║  Shared Memory Block Allocator                                               ║
╠══════════════════════════════════════════════════════════════════════════════╣
║  Version : 2.0.1                                                             ║
║  Date    : 2026-02-07                                                        ║
╠══════════════════════════════════════════════════════════════════════════════╣
║  Summary:                                                                    ║
║    프로세스 간 공유 메모리 이미지 블록 풀 관리자                             ║
║    - Free List 알고리즘 O(1) 블록 할당/해제                                  ║
║    - Layout: [Index: int16[N]] + [Data: uint8[N][H][W][C]] + [Free Head]     ║
╠══════════════════════════════════════════════════════════════════════════════╣
║  Classes:                                                                    ║
║    SmBlock        - 공유 메모리 블록 풀                                      ║
║    SmBlockHandler - SmBlock 생성/주입 핸들러                                 ║
╚══════════════════════════════════════════════════════════════════════════════╝
"""

from multiprocessing.shared_memory import SharedMemory
from typing import Tuple, Optional, Dict, Any
import numpy as np


class SmBlock:
    """프로세스 간 공유 메모리 이미지 블록 풀 관리자."""

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

    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

        # 총 공유 메모리 크기
        total_size = self._index_size + (self._item_size * maxsize) + 2

        # 공유 메모리 생성 또는 연결
        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]
        )

        # Free Head: offset 2*N + N*H*W*C, size 2
        free_head_offset = data_offset + data_size
        self._free_head = np.ndarray(
            (1,), dtype=np.int16,
            buffer=buf[free_head_offset:free_head_offset + 2]
        )

    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 블록

    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}")

        with self._lock:
            head = self._free_head[0]
            if head == -1:
                return -1  # 풀이 가득 참

            # Free list에서 블록 분리
            next_free = self._index_buffer[head]
            self._free_head[0] = next_free
            self._index_buffer[head] = -1  # 사용 중 표시

            # 데이터 복사
            self._data_buffer[head][:] = image
            self._last_alloc = head
            return int(head)

    def malloc2(self, image: np.ndarray) -> int:
        """마지막 할당 위치 다음부터 순차 탐색하여 빈 공간 할당.

        Args:
            image: 저장할 이미지

        Returns:
            할당된 블록 인덱스 또는 -1
        """
        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]:
                    # 빈 블록 발견 - Free list에서 제거
                    if idx == self._free_head[0]:
                        # head인 경우
                        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._data_buffer[idx][:] = image
                    self._last_alloc = idx
                    return idx

            return -1  # 풀이 가득 참

    def alloc(self) -> int:
        """인덱스만 할당 (데이터 복사 없음). O(1) 시간복잡도.

        Zero-copy 패턴: alloc() → get_buffer() → 직접 쓰기

        Returns:
            할당된 블록 인덱스, 풀이 가득 차면 -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._last_alloc = head
            return int(head)

    def get_buffer(self, index: int) -> np.ndarray:
        """할당된 블록의 쓰기 가능한 버퍼 참조 반환 (복사 없음).

        Args:
            index: 블록 인덱스

        Returns:
            공유 메모리 버퍼 참조 (쓰기 가능)
        """
        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에 반환.

        Args:
            index: 해제할 블록 인덱스

        Raises:
            ValueError: 이미 비어있는 블록을 해제할 때
            IndexError: 유효하지 않은 인덱스
        """
        if index < 0 or index >= self._maxsize:
            raise IndexError(f"Index {index} out of range [0, {self._maxsize})")

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

            # Free list의 head로 추가 (스택 방식)
            self._index_buffer[index] = self._free_head[0]
            self._free_head[0] = index

    def get(self, index: int) -> np.ndarray:
        """할당된 블록의 이미지 데이터 반환.

        Args:
            index: 읽을 블록 인덱스

        Returns:
            이미지 데이터 (shape=생성자의 shape)

        Raises:
            ValueError: 비어있는 블록에 접근할 때
            IndexError: 유효하지 않은 인덱스
        """
        if index < 0 or index >= self._maxsize:
            raise IndexError(f"Index {index} out of range [0, {self._maxsize})")

        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:
        """현재 사용 가능한(비어있는) 블록 수 반환."""
        with self._lock:
            cnt = 0
            idx = self._free_head[0]
            while idx != -1:
                cnt += 1
                idx = self._index_buffer[idx]
            return cnt

    def close(self):
        """공유 메모리 연결 해제 및 정리."""
        self._shm.close()
        if self._is_new:
            self._shm.unlink()

    @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 __str__(self) -> str:
        return f"SmBlock[ {self._name}, {self._shape}, {self._maxsize} ]"

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


class SmBlockHandler:
    """SmBlock 생성/주입 핸들러. TaskManager에서 사용."""
    _pools: Dict[str, 'SmBlock'] = {}

    @classmethod
    def creation(cls, gconfig, mgr, tasks: Dict) -> Dict[str, 'SmBlock']:
        """_smblock 설정으로 SmBlock 생성 및 injection dict 변환.

        Args:
            gconfig: GConfig instance containing:
                - platform_config._smblock: (new format) Shared memory block definitions
                - task_config._smblock: (old format, fallback) Shared memory block definitions
        """
        # New format: platform_config._smblock, fallback: task_config._smblock
        platform_config = gconfig.data_get("platform_config", {})
        task_config = gconfig.data_get("task_config", {})
        smblock_cfg = platform_config.get("_smblock", {}) or task_config.get("_smblock", {})
        if not smblock_cfg:
            return cls._pools

        smblock_info = {}
        for name, cfg in smblock_cfg.items():
            if name in cls._pools:
                continue
            shape = cfg.get("shape")
            maxsize = cfg.get("maxsize", 100)
            if not shape:
                print(f"[SmBlockHandler] '{name}': shape required")
                continue
            shape = tuple(shape) if isinstance(shape, list) else shape
            lock = mgr.Lock()
            try:
                pool = SmBlock(name, shape, maxsize, create=True, lock=lock)
                cls._pools[name] = pool
                smblock_info[name] = {"name": name, "shape": shape, "maxsize": maxsize, "lock": lock}
            except FileExistsError:
                from multiprocessing.shared_memory import SharedMemory
                try:
                    old = SharedMemory(name=name, create=False)
                    old.close()
                    old.unlink()
                except Exception:
                    pass
                pool = SmBlock(name, shape, maxsize, create=True, lock=lock)
                cls._pools[name] = pool
                smblock_info[name] = {"name": name, "shape": shape, "maxsize": maxsize, "lock": lock}
            except Exception as e:
                print(f"[SmBlockHandler] '{name}' failed: {e}")

        # injection 문자열 → dict 변환
        for tid, ti in tasks.items():
            for k, v in list(ti.injection.items()):
                if isinstance(v, str) and v.startswith("smblock:"):
                    pool_name = v[8:]
                    if pool_name in smblock_info:
                        ti.injection[k] = {"_smblock": pool_name, **smblock_info[pool_name]}
                    else:
                        available = list(smblock_info.keys()) or "(none)"
                        raise ValueError(f"[SmBlockHandler] {tid}.{k}: '{pool_name}' not defined. Available: {available}")

        # 테이블 출력
        if cls._pools:
            cls._print_smblock_table()
        return cls._pools

    @classmethod
    def _print_smblock_table(cls):
        """SmBlock 정보 테이블 출력."""
        headers = ['Name', 'Shape', 'MaxSize', 'Size(MB)']
        rows = []
        for name, pool in cls._pools.items():
            shape_str = str(pool.shape)
            total_bytes = pool.index_size + (pool.item_size * pool.maxsize) + 2
            size_mb = f"{total_bytes / (1024 * 1024):.1f}"
            rows.append([name, shape_str, str(pool.maxsize), size_mb])
        widths = [max(len(h), max((len(r[i]) for r in rows), default=0)) for i, h in enumerate(headers)]
        top    = '┌' + '┬'.join('─' * (w + 2) for w in widths) + '┐'
        mid    = '├' + '┼'.join('─' * (w + 2) for w in widths) + '┤'
        bottom = '└' + '┴'.join('─' * (w + 2) for w in widths) + '┘'
        print('\n' + top)
        print('│' + '│'.join(f' {h:<{w}} ' for h, w in zip(headers, widths)) + '│')
        print(mid)
        for row in rows:
            print('│' + '│'.join(f' {c:<{w}} ' for c, w in zip(row, widths)) + '│')
        print(bottom + '\n')

    @classmethod
    def injection(cls, worker, _tasks: Dict):
        """injection dict를 SmBlock 인스턴스로 변환하여 job에 주입."""
        job, ti = worker.job, worker.ti
        for k, v in ti.injection.items():
            if isinstance(v, dict) and "_smblock" in v:
                pool_name = v["_smblock"]
                shape = tuple(v["shape"])
                maxsize = v["maxsize"]
                lock = v["lock"]
                if pool_name in cls._pools:
                    setattr(job, k, cls._pools[pool_name])
                else:
                    try:
                        pool = SmBlock(pool_name, shape, maxsize, create=False, lock=lock)
                        setattr(job, k, pool)
                    except Exception as e:
                        print(f"[SmBlockHandler] {ti.task_id}.{k}: attach failed: {e}")

    @classmethod
    def close_all(cls):
        """모든 SmBlock 정리."""
        for name, pool in cls._pools.items():
            try:
                pool.close()
            except Exception as e:
                print(f"[SmBlockHandler] '{name}' close failed: {e}")
        cls._pools.clear()
