openclaw/agents/stock/workspace/sim/universe.py

"""종목 풀 조립 — sim 전용 누적 관찰목록(state/sim/watchlist.json).

자산 소스(비하이브 워치/관심 + 본인 보유)에 새 종목이 생기면 자동 편입하되,
한 번 들어온 종목은 자동으로 빠지지 않는다(삭제는 수동만). 관심종목 탭에서 직접 추가/삭제 가능.

각 종목에 origin 태깅: watch / interest / held / manual. 가격선(비하이브 buy/target/stop)은
의도적으로 무시한다 — 진입·손절·목표는 sim 엔진이 데이터로 직접 계산한다.
"""
from __future__ import annotations

import json
import sys
from datetime import datetime

from . import config

sys.path.insert(0, str(config.SCRIPTS))
import kiwoom_client as kc  # noqa: E402

WATCHLIST_PATH = config.WORKSPACE / 'state' / 'behive_watchlist.json'
INTERESTS_PATH = config.WORKSPACE / 'state' / 'behive_interests.json'
SIM_WATCHLIST_PATH = config.STATE_DIR / 'watchlist.json'  # sim 누적 관찰목록


def _load_json(path, default):
    try:
        return json.loads(path.read_text())
    except Exception:
        return default


def _load_watchlist() -> dict:
    """sim 누적 관찰목록 {code: {name, origin, added_at}}."""
    wl = _load_json(SIM_WATCHLIST_PATH, {})
    return wl if isinstance(wl, dict) else {}


def _save_watchlist(wl: dict):
    config.STATE_DIR.mkdir(parents=True, exist_ok=True)
    SIM_WATCHLIST_PATH.write_text(json.dumps(wl, ensure_ascii=False, indent=2))


def _collect_sources() -> dict[str, tuple[str, str]]:
    """자산 소스(비하이브 워치/관심 + 본인 보유)에서 {code: (name, origin)} 수집. 코드 6자리 정규화."""
    out: dict[str, tuple[str, str]] = {}

    def add(code: str, name: str, origin: str):
        code = (code or '').strip().zfill(6) if code else ''
        if not code or code in out:
            return
        out[code] = (name or '', origin)

    watch = _load_json(WATCHLIST_PATH, {})
    for name, info in (watch.items() if isinstance(watch, dict) else []):
        add(info.get('code', ''), info.get('stock', name), 'watch')

    interests = _load_json(INTERESTS_PATH, {})
    for name, info in (interests.items() if isinstance(interests, dict) else []):
        add(info.get('code', ''), info.get('stock', name), 'interest')

    try:
        held = kc.get_positions_all(labels=config.OWNER_ACCOUNT_LABELS)
        for _label, positions in held.items():
            for p in positions:
                if p.get('qty', 0) > 0:
                    add(p.get('code', ''), p.get('name', ''), 'held')
    except Exception as e:
        sys.stderr.write(f'[universe] 보유종목 조회 실패 (무시): {e}\n')

    return out


def sync_watchlist() -> dict:
    """자산 소스의 새 종목만 누적 관찰목록에 편입(이름 보강). 자동 삭제는 안 함. 갱신된 목록 반환."""
    wl = _load_watchlist()
    now = datetime.now(config.KST).isoformat()
    changed = False
    for code, (name, origin) in _collect_sources().items():
        ent = wl.get(code)
        if ent is None:
            wl[code] = {'name': name, 'origin': origin, 'added_at': now}
            changed = True
        elif not ent.get('name') and name:  # 이름만 보강
            ent['name'] = name
            changed = True
    if changed:
        _save_watchlist(wl)
    return wl


def search_stocks(query: str, limit: int = 20) -> list[dict]:
    """종목명/코드 부분검색 → [{code, name}] 후보 (자동완성 팝업용). ci-exact > prefix > 부분 정렬."""
    q = (query or '').strip()
    if not q:
        return []
    cache = kc._load_code_cache()
    if q.isdigit() and len(q) == 6:  # 6자리 코드 직접
        try:
            info = kc.resolve_stock_code(q)
            return [{'code': info.get('code', q), 'name': info.get('name', '')}]
        except Exception:
            return []
    ql = q.lower()
    hits = [(name, info) for name, info in cache.items() if ql in (name or '').lower()]

    def sk(item):
        n = (item[0] or '').lower()
        rank = 0 if n == ql else 1 if n.startswith(ql) else 2
        return (rank, len(n), n)  # 같은 등급이면 짧은 이름(핵심 종목) 우선 → '삼성전자'가 ETN보다 위

    hits.sort(key=sk)
    return [{'code': info.get('code', ''), 'name': info.get('name') or name}
            for name, info in hits[:limit] if info.get('code')]


def add_manual(code_or_name: str) -> dict | None:
    """관심종목 탭 수동 추가 — 코드/이름 해석 후 origin='manual'로 편입. 이미 있으면 기존 반환."""
    try:
        info = kc.resolve_stock_code(code_or_name)
    except Exception as e:
        sys.stderr.write(f'[universe] resolve 실패: {e}\n')
        return None
    code = (info.get('code') or '').strip().zfill(6)
    if not code:
        return None
    wl = _load_watchlist()
    if code not in wl:
        wl[code] = {'name': info.get('name') or '', 'origin': 'manual',
                    'added_at': datetime.now(config.KST).isoformat()}
        _save_watchlist(wl)
    return {'code': code, **wl[code]}


def remove(code: str):
    """관심종목 탭 수동 삭제."""
    code = (code or '').strip().zfill(6)
    wl = _load_watchlist()
    if code in wl:
        del wl[code]
        _save_watchlist(wl)


def build_universe() -> list[dict]:
    """누적 관찰목록을 동기화 후 [{code, name, sources:[origin]}]로 반환. code 6자리 정규화."""
    wl = sync_watchlist()
    return [{'code': code, 'name': ent.get('name', ''), 'sources': [ent.get('origin', 'manual')]}
            for code, ent in wl.items()]


def code_market_map() -> dict[str, str]:
    """code → 'KOSPI'/'KOSDAQ' 매핑 (종목코드 캐시 기반). 모르면 빈값."""
    cache = kc._load_code_cache()
    out: dict[str, str] = {}
    for _name, meta in cache.items():
        code = (meta.get('code') or '').strip()
        if code:
            out[code.zfill(6)] = meta.get('market') or ''
    return out


if __name__ == '__main__':
    u = build_universe()
    print(f'universe: {len(u)} 종목')
    for e in u:
        print(f"  {e['code']} {e['name']:12s} {'/'.join(e['sources'])}")