#!/usr/bin/env python3
"""
kis_trader/backtest/shared_ticks.py — 파라서치 ws_ticks 공유메모리 저장 (spawn-safe)
=====================================================================================
[왜 필요한가]
파라서치 워커는 spawn(= max_tasks_per_child 가 강제) 방식이라, 워커마다 틱 dict
(관측 ~3.3GB) 을 각자 디스크 pickle 에서 재역직렬화한다 → 메모리 N배·기동 지연.
fork COW 로 공유하려 하면 부모가 멀티스레드(DB·로깅)라 fork+threads 데드락이 난다.

[해결]
틱을 컬럼(numpy 배열) 로 평탄화해 ``multiprocessing.shared_memory`` 에 **1벌**만 올린다.
워커는 블록 '이름'으로 read-only attach(mmap → 물리 사본 1개 공유). 접근 시 해당 분(minute)
틱만 즉석에서 dict 로 만들어 반환한다(원본 스키마 100% 동일). fork 가 아니라 spawn-safe.

[원본 자료구조]
    {code: {minute_key(12자리): [ {code,tick_time,price,volume,source}, ... ]}}
[공유 표현]
    평탄화 배열 price(f8) / volume(i8) / tick_time(S14) / source(S..) / epoch(i8) + 인덱스
    {code: {minute: (offset, count)}}
    epoch(=_tick_epoch_sec) 은 사전계산해, 핫루프(휩쏘/청산)가 문자열을 재파싱하지 않고
    ``TickColumnView`` 로 배열을 직접 읽게 한다(dict 재구성 제거 → 속도 회복).

[동일성 보장]
읽기 래퍼(SharedTicksMapping / SharedBucketMapping)는 ``collections.abc.Mapping`` 이라
dict 프로토콜(.get/.keys/.values/.items/in/len/bool/dict()) 을 원본 dict 와 동일하게 제공한다.
분(minute) 값은 lazy 시퀀스(_TickList)로 감싸, ``len()``·``any()``·``bool()`` 같은
콤보별 집계 스캔은 dict 를 만들지 않고 O(1) 로 답하고, 실제로 순회(list()/extend())할 때만
원본과 완전히 같은 tick dict 를 생성한다.

기본은 사용하지 않는다(호출부에서 env 플래그로 opt-in). numpy/shared_memory 가 없으면
build_shared_ticks 가 None 을 반환해 호출부가 기존(디스크 pickle) 경로로 자연히 폴백한다.
"""
from __future__ import annotations

from collections.abc import Mapping, Sequence
from typing import Any, Dict, List, Optional, Tuple

try:
    import numpy as np
    from multiprocessing import shared_memory
    _SHM_AVAILABLE = True
except Exception:  # pragma: no cover - numpy 미설치 환경 폴백
    _SHM_AVAILABLE = False


def shared_ticks_available() -> bool:
    """numpy + shared_memory 사용 가능 여부."""
    return _SHM_AVAILABLE


# 컬럼 dtype (엔디안 고정 — 다른 프로세스에서 attach 해도 동일 해석)
_DT_PRICE = "<f8"
_DT_VOLUME = "<i8"
_DT_TICKTIME = "S14"  # tick_time 은 로더에서 [:14] → 최대 14바이트 (12자리도 널패딩→디코드시 복원)
_DT_EPOCH = "<i8"  # tick_time → epoch초 사전계산 (핫루프 문자열 파싱 제거, whipsaw._tick_epoch_sec 동일)


def _unregister_shm_from_tracker(name: str) -> None:
    """워커(attach 전용)가 종료할 때 shared_memory 를 unlink 하지 못하게 한다.

    파이썬 resource_tracker 는 attach 만 한 자식도 종료 시 소유자로 오인해 unlink 를
    시도하는 알려진 함정이 있다(부모가 아직 쓰는데 블록이 사라짐). attach 직후 등록을
    해제해, unlink 권한은 오직 '생성한 부모'에게만 남긴다.
    """
    try:
        from multiprocessing import resource_tracker  # type: ignore
        resource_tracker.unregister(f"/{name}", "shared_memory")
    except Exception:
        pass


class _TickList(Sequence):
    """한 분(minute) 의 틱 리스트 — 순회할 때만 원본과 동일한 tick dict 를 생성.

    ``len()`` / ``bool()`` 은 O(1)(count 만) 이라, 콤보마다 도는 커버리지·집계 스캔이
    dict 를 만들지 않고 값싸게 끝난다. ``list(...)`` / ``extend(...)`` / 인덱싱 시에만
    실제 dict 를 만든다.
    """

    __slots__ = ("_owner", "_code", "_off", "_cnt")

    def __init__(self, owner: "SharedTicksMapping", code: str, off: int, cnt: int) -> None:
        self._owner = owner
        self._code = code
        self._off = off
        self._cnt = cnt

    def __len__(self) -> int:
        return self._cnt

    def __getitem__(self, i):
        if isinstance(i, slice):
            start, stop, step = i.indices(self._cnt)
            return [self._owner._make_tick(self._code, self._off + j)
                    for j in range(start, stop, step)]
        if i < 0:
            i += self._cnt
        if i < 0 or i >= self._cnt:
            raise IndexError(i)
        return self._owner._make_tick(self._code, self._off + i)

    def column_view(self) -> "TickColumnView":
        """이 분(minute) 구간을 dict 재구성 없이 배열 직접 읽는 컬럼 뷰로 반환.

        ``bucket.get(minute)`` 가 반환한 _TickList 를 그대로 컬럼 뷰(단일 세그먼트)로
        바꿔, 백테 틱재생 핫루프(breakout intrabar 등)가 dict 생성 없이 순회하도록 한다.
        저장 순서(로더 ORDER BY tick_time)를 그대로 보존 → dict 경로 정렬과 동일.
        """
        self._owner._ensure()
        if self._cnt <= 0:
            return TickColumnView(self._owner, [])
        return TickColumnView(self._owner, [(self._off, self._cnt)])


class TickColumnView:
    """공유 컬럼 배열 위의 '틱 구간 뷰' — dict 재구성 없이 배열을 직접 읽는다.

    [왜]
    분(minute) 값을 dict 리스트로 재구성하면(_TickList) 틱마다 문자열 디코드·형변환
    비용이 들어, 엔진이 같은 구간을 반복 스캔하는 휩쏘/청산 핫루프에서 dict판 대비
    ~100배 느려진다(실측). 이 뷰는 owner 의 컬럼 배열(_epoch/_price/_volume/_tick_time)을
    절대 인덱스로 직접 읽어 그 재구성 비용을 없앤다.

    segments: [(offset, count), ...] — 원본 merged 순서를 그대로 보존한다.
    ``iter_idx()`` 는 그 순서대로 절대 인덱스를 산출하므로, dict 경로에서
    ``merged.extend(bucket[mk])`` 로 만든 리스트 순회와 완전히 동일한 순서다.
    """

    __slots__ = ("owner", "segments", "_len")

    def __init__(self, owner: "SharedTicksMapping", segments: List[Tuple[int, int]]) -> None:
        self.owner = owner
        self.segments = segments
        total = 0
        for _off, cnt in segments:
            total += cnt
        self._len = total

    def __len__(self) -> int:
        return self._len

    def __bool__(self) -> bool:
        return self._len > 0

    def iter_idx(self):
        """merged 순서대로 절대 배열 인덱스 산출 (dict 리스트 순회와 동일 순서)."""
        for off, cnt in self.segments:
            end = off + cnt
            i = off
            while i < end:
                yield i
                i += 1

    def cap_by_tick_time_le(self, key14: str) -> "TickColumnView":
        """tick_time[:14] <= key14 인 틱만 남긴 새 뷰 (라이브스캔큐 매도 캡핑용).

        dict 경로 ``[tk for tk in minute_ticks if str(tk['tick_time'])[:14] <= scan_key[:14]]``
        와 동일 semantics — 매 틱 문자열([:14]) 비교로 통과분을 순서대로 모아 세그먼트로
        합친다(정렬 가정에 의존하지 않아, 원본 리스트 컴프리헨션과 bit-identical).
        """
        ttarr = self.owner._tick_time
        segs: List[Tuple[int, int]] = []
        cur_off = -1
        cur_cnt = 0
        for i in self.iter_idx():
            tt = ttarr[i].decode("utf-8")
            if tt[:14] <= key14:
                if cur_off < 0:
                    cur_off, cur_cnt = i, 1
                elif i == cur_off + cur_cnt:
                    cur_cnt += 1
                else:
                    segs.append((cur_off, cur_cnt))
                    cur_off, cur_cnt = i, 1
        if cur_off >= 0:
            segs.append((cur_off, cur_cnt))
        return TickColumnView(self.owner, segs)

    def first_idx(self) -> int:
        for off, cnt in self.segments:
            if cnt > 0:
                return off
        return -1

    def last_idx(self) -> int:
        for off, cnt in reversed(self.segments):
            if cnt > 0:
                return off + cnt - 1
        return -1


class SharedBucketMapping(Mapping):
    """한 종목의 {minute_key: [tick,...]} 매핑 — 값은 lazy _TickList.

    Mapping 상속으로 .get/.keys/.values/.items/in/len/bool/dict() 를 원본 dict 와
    동일하게 제공한다(누락 키 → .get 은 None, 원본 동작과 일치).
    """

    __slots__ = ("_owner", "_code", "_cidx")

    def __init__(self, owner: "SharedTicksMapping", code: str,
                 cidx: Dict[str, Tuple[int, int]]) -> None:
        self._owner = owner
        self._code = code
        self._cidx = cidx

    def __getitem__(self, minute: str) -> _TickList:
        off, cnt = self._cidx[minute]  # KeyError → Mapping.get 이 기본값 처리
        self._owner._ensure()
        return _TickList(self._owner, self._code, off, cnt)

    def __iter__(self):
        return iter(self._cidx)

    def __len__(self) -> int:
        return len(self._cidx)

    # ── 컬럼 직접접근 뷰 (백테 틱재생 핫루프 전용) ───────────────────────
    def column_view_minute(self, minute_key: str) -> "TickColumnView":
        """단일 분(minute) 구간 뷰 — collect_minute_ticks 컬럼판.

        원본은 해당 분 틱을 tick_time 오름차순 정렬해 반환하는데, 로더가 이미
        ``ORDER BY code, tick_time`` 로 담아둔 저장 순서가 오름차순이라 정렬은
        무연산(stable no-op)이다 → 저장 순서 뷰가 정렬 결과와 완전히 동일하다.
        """
        self._owner._ensure()
        seg = self._cidx.get(str(minute_key)[:12])
        if seg is None or seg[1] <= 0:
            return TickColumnView(self._owner, [])
        return TickColumnView(self._owner, [seg])

    def column_view_lookback(self, minute_keys: List[str]) -> "TickColumnView":
        """여러 분(minute) 을 주어진 순서대로 이어붙인 뷰 — collect_whipsaw_ticks 컬럼판.

        minute_keys 는 호출부에서 이미 원하는 순서(과거→최신)로 정렬해 넘긴다.
        원본 ``merged.extend(bucket.get(mk) or [])`` 와 동일하게, 없는/빈 분은 건너뛴다.
        """
        self._owner._ensure()
        segs: List[Tuple[int, int]] = []
        for mk in minute_keys:
            seg = self._cidx.get(mk)
            if seg is not None and seg[1] > 0:
                segs.append(seg)
        return TickColumnView(self._owner, segs)


class SharedTicksMapping(Mapping):
    """워커측 최상위 {code: bucket} 매핑 — 원본 ticks_by_code 를 그대로 대체.

    descriptor(이름·dtype·인덱스) 만 있으면 어느 프로세스에서든 attach 해 쓴다.
    실제 shared_memory attach 는 처음 종목 접근 시 1회(lazy).
    """

    def __init__(self, descriptor: Dict[str, Any]) -> None:
        self._meta = descriptor
        self._index: Dict[str, Dict[str, Tuple[int, int]]] = descriptor["index"]
        self._total: int = int(descriptor["total"])
        self._src_dtype: str = descriptor["src_dtype"]
        self._names: Dict[str, str] = descriptor["names"]
        self._attached = False
        self._shms: Dict[str, Any] = {}
        self._price = None
        self._volume = None
        self._tick_time = None
        self._source = None
        self._epoch = None

    # ── attach (워커 1회) ────────────────────────────────────────────────
    def _ensure(self) -> None:
        if self._attached:
            return
        for key, name in self._names.items():
            shm = shared_memory.SharedMemory(name=name)
            self._shms[key] = shm
            _unregister_shm_from_tracker(name)  # attach 전용 → unlink 권한 포기
        self._price = np.ndarray((self._total,), dtype=_DT_PRICE, buffer=self._shms["price"].buf)
        self._volume = np.ndarray((self._total,), dtype=_DT_VOLUME, buffer=self._shms["volume"].buf)
        self._tick_time = np.ndarray((self._total,), dtype=_DT_TICKTIME, buffer=self._shms["tick_time"].buf)
        self._source = np.ndarray((self._total,), dtype=self._src_dtype, buffer=self._shms["source"].buf)
        self._epoch = np.ndarray((self._total,), dtype=_DT_EPOCH, buffer=self._shms["epoch"].buf)
        self._attached = True

    def _make_tick(self, code: str, i: int) -> Dict[str, Any]:
        """원본 로더가 만든 tick dict 와 100% 동일한 dict 를 복원."""
        return {
            "code": code,
            "tick_time": self._tick_time[i].decode("utf-8"),
            "price": float(self._price[i]),
            "volume": int(self._volume[i]),
            "source": self._source[i].decode("utf-8"),
        }

    def __getitem__(self, code: str) -> SharedBucketMapping:
        cidx = self._index[code]  # KeyError → Mapping.get 기본값 처리
        self._ensure()
        return SharedBucketMapping(self, code, cidx)

    def __iter__(self):
        return iter(self._index)

    def __len__(self) -> int:
        return len(self._index)

    def close(self) -> None:
        """워커/사용측 attach 해제(블록 삭제 아님 — unlink 는 부모만)."""
        for shm in self._shms.values():
            try:
                shm.close()
            except Exception:
                pass
        self._shms = {}
        self._attached = False


class SharedTickStore:
    """부모: 컬럼 배열을 shared_memory 에 올리고 핸들을 보관(=블록 생존 유지).

    - descriptor(): 워커로 넘길 가벼운 dict (shm 이름·dtype·shape + 인덱스).
    - unlink(): 탐색 종료 후 블록 삭제(부모만 호출).
    """

    def __init__(self, shms: Dict[str, Any], meta: Dict[str, Any]) -> None:
        self._shms = shms
        self._meta = meta

    def descriptor(self) -> Dict[str, Any]:
        return self._meta

    def close(self) -> None:
        for shm in self._shms.values():
            try:
                shm.close()
            except Exception:
                pass

    def unlink(self) -> None:
        for shm in self._shms.values():
            try:
                shm.close()
            except Exception:
                pass
            try:
                shm.unlink()
            except Exception:
                pass
        self._shms = {}

    def attach_mapping(self) -> SharedTicksMapping:
        """같은 프로세스(부모)에서도 검증용으로 매핑 뷰를 얻는다(이름으로 재-attach)."""
        return SharedTicksMapping(self._meta)

    def local_mapping(self) -> SharedTicksMapping:
        """같은 프로세스(Optuna 단일프로세스/스레드)용 — 부모가 이미 연 버퍼를 그대로 재사용.

        cross-process 용 ``attach_mapping`` 은 워커에서 ``SharedMemory(name=...)`` 로 다시
        attach 하지만, 같은 프로세스에서 그렇게 하면 resource_tracker 재등록/해제가 엇갈려
        종료 시 무해하지만 시끄러운 ``KeyError`` 노이즈가 난다. 여기서는 새 attach 없이
        생성자가 보관 중인 shm 버퍼로 numpy 배열을 직접 얹어(추가 attach 0회) 그 노이즈를
        원천 제거한다. 배열을 미리 연결하므로 스레드(n_jobs>1) 초기화 레이스도 없다.
        """
        m = SharedTicksMapping(self._meta)
        total = int(self._meta["total"])
        src_dtype = self._meta["src_dtype"]
        m._price = np.ndarray((total,), dtype=_DT_PRICE, buffer=self._shms["price"].buf)
        m._volume = np.ndarray((total,), dtype=_DT_VOLUME, buffer=self._shms["volume"].buf)
        m._tick_time = np.ndarray((total,), dtype=_DT_TICKTIME, buffer=self._shms["tick_time"].buf)
        m._source = np.ndarray((total,), dtype=src_dtype, buffer=self._shms["source"].buf)
        m._epoch = np.ndarray((total,), dtype=_DT_EPOCH, buffer=self._shms["epoch"].buf)
        m._attached = True  # 배열 이미 연결 → _ensure() no-op (재-attach 안 함)
        return m


def build_shared_ticks(
    ticks_by_code: Optional[Dict[str, Dict[str, List[Dict[str, Any]]]]],
) -> Optional[SharedTickStore]:
    """원본 ticks_by_code 를 공유메모리 컬럼 저장으로 변환.

    실패/빈 데이터/미지원 환경이면 None → 호출부가 기존 경로로 폴백.
    평탄화 순서는 code(정렬)→minute(원본 insertion order)→기존 리스트 순 으로 고정한다.
    """
    if not _SHM_AVAILABLE or not ticks_by_code:
        return None

    # epoch 사전계산은 실매/백테 공용 _tick_epoch_sec 와 100% 동일해야 한다 →
    # 복제하지 않고 그 함수를 그대로 import (import 실패 시 폴백 = 디스크 경로).
    try:
        from kis_trader.engine.whipsaw_filter import _tick_epoch_sec
    except Exception:
        return None

    codes = sorted(ticks_by_code.keys())

    # 1) 총 틱 수 + source 최대 바이트 길이 (S dtype 폭 결정)
    total = 0
    max_src = 1
    for code in codes:
        bucket = ticks_by_code.get(code) or {}
        for mk in bucket:
            lst = bucket.get(mk) or []
            total += len(lst)
            for t in lst:
                s = str(t.get("source") or "")
                b = len(s.encode("utf-8"))
                if b > max_src:
                    max_src = b
    if total <= 0:
        return None
    src_dtype = f"S{max_src}"

    specs = {
        "price": np.dtype(_DT_PRICE),
        "volume": np.dtype(_DT_VOLUME),
        "tick_time": np.dtype(_DT_TICKTIME),
        "source": np.dtype(src_dtype),
        "epoch": np.dtype(_DT_EPOCH),
    }
    shms: Dict[str, Any] = {}
    arrs: Dict[str, Any] = {}
    try:
        for name, dt in specs.items():
            nbytes = int(dt.itemsize) * total
            shm = shared_memory.SharedMemory(create=True, size=max(1, nbytes))
            shms[name] = shm
            arrs[name] = np.ndarray((total,), dtype=dt, buffer=shm.buf)
    except Exception:
        # 생성 도중 실패 → 이미 만든 블록 정리 후 폴백
        for shm in shms.values():
            try:
                shm.close(); shm.unlink()
            except Exception:
                pass
        return None

    price = arrs["price"]
    volume = arrs["volume"]
    tt = arrs["tick_time"]
    src = arrs["source"]
    epoch = arrs["epoch"]

    index: Dict[str, Dict[str, Tuple[int, int]]] = {}
    off = 0
    for code in codes:
        bucket = ticks_by_code.get(code) or {}
        cidx: Dict[str, Tuple[int, int]] = {}
        for mk in bucket:  # 원본 dict insertion order 보존 (iteration 동일성)
            lst = bucket.get(mk) or []
            cnt = len(lst)
            start = off
            for t in lst:
                tt_str = str(t.get("tick_time") or "")
                price[off] = float(t.get("price") or 0)
                volume[off] = int(t.get("volume") or 0)
                tt[off] = tt_str[:14].encode("utf-8")
                src[off] = str(t.get("source") or "").encode("utf-8")
                epoch[off] = _tick_epoch_sec(tt_str)  # 핫루프에서 재파싱 안 하도록 사전계산
                off += 1
            cidx[mk] = (start, cnt)
        index[code] = cidx

    meta = {
        "total": total,
        "src_dtype": src_dtype,
        "names": {k: shms[k].name for k in specs},
        "index": index,
    }
    return SharedTickStore(shms, meta)


def build_shared_ticks_view(
    ticks_by_code: Optional[Dict[str, Dict[str, List[Dict[str, Any]]]]],
    *,
    enabled: bool = True,
) -> Tuple[Any, Optional["SharedTickStore"]]:
    """단일 프로세스(Optuna 등)용 — dict 를 컬럼 공유메모리 '뷰' 로 치환해 RAM 절감.

    [Grid 와 차이]
      Grid 파라서치는 워커 '프로세스' 간 사본 제거가 목적이라 descriptor 를 워커로
      넘겨 attach 시킨다. Optuna 는 단일 프로세스(또는 n_jobs 스레드)라 사본 문제는
      없지만, dict-of-dict(파이썬 객체 오버헤드 큼)을 numpy 컬럼(shared_memory)으로
      바꾸면 **메모리 사용량**이 크게 준다 → RAM 이 적은 WSL 에서도 대용량 틱으로
      Optuna 를 돌릴 수 있다.

    반환 ``(view, store)``:
      · 성공: ``(SharedTicksMapping, SharedTickStore)`` — 종료 시 ``store.unlink()`` 필수
      · 미지원/빈데이터/off/실패: ``(원본 ticks_by_code, None)`` — 호출부가 dict 경로 그대로 사용

    ``SharedTicksMapping`` 은 원본 dict 와 bit-identical(Grid E2E 검증) 이라 evaluate 무변경.
    """
    if not enabled or not ticks_by_code or not shared_ticks_available():
        return ticks_by_code, None
    store = build_shared_ticks(ticks_by_code)
    if store is None:
        return ticks_by_code, None
    # 단일 프로세스용 로컬 매핑 — 부모 버퍼 재사용(추가 attach 0회) → 종료 noise 없음 + 스레드 안전.
    view = store.local_mapping()
    return view, store


__all__ = [
    "shared_ticks_available",
    "build_shared_ticks",
    "build_shared_ticks_view",
    "SharedTickStore",
    "SharedTicksMapping",
    "SharedBucketMapping",
    "TickColumnView",
]