#!/usr/bin/env python3
import argparse
import time
import socket
import json

# ~~~~~============== CONFIGURATION  ==============~~~~~
team_name = "HanyangFloorFunction"

# --- 상태 관리 및 유틸리티 ---
class Dir:
    BUY = "BUY"
    SELL = "SELL"

class StateManager:
    def __init__(self):
        self.positions = {}
        self.bid_prices = {}
        self.ask_prices = {}

    def update_position(self, symbol, delta):
        self.positions[symbol] = self.positions.get(symbol, 0) + delta

    def get_position(self, symbol):
        return self.positions.get(symbol, 0)

    def update_bid_ask_price(self, symbol, bid, ask):
        self.bid_prices[symbol] = bid
        self.ask_prices[symbol] = ask

class OrderManager:
    def __init__(self):
        self.order_id_counter = 0
    def next_order(self):
        self.order_id_counter += 1
        return self.order_id_counter

# --- 연결 클래스 ---
class ExchangeConnection:
    def __init__(self, args):
        self.exchange_hostname = args.exchange_hostname
        self.port = args.port
        # 타임아웃 설정으로 무한 대기 방지
        self.s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.s.settimeout(5)
        self.s.connect((self.exchange_hostname, self.port))
        self.reader = self.s.makefile("r", 1)
        # 초기 Handshake
        self._write_message({"type": "hello", "team": team_name.upper()})

    def read_message(self):
        line = self.reader.readline()
        if not line: return None
        return json.loads(line)

    def _write_message(self, message):
        # 모든 메시지 끝에는 줄바꿈(\n)이 필수입니다
        self.s.send((json.dumps(message) + "\n").encode("utf-8"))

    def send_add_message(self, order_id, symbol, direction, price, size):
        self._write_message({
            "type": "add", "order_id": order_id, "symbol": symbol, 
            "dir": direction, "price": price, "size": size, "tif": "DAY"
        })

# --- 메인 봇 로직 ---
def main():
    args = parse_arguments()
    try:
        exchange = ExchangeConnection(args)
    except Exception as e:
        print(f"연결 실패: {e}")
        return

    state = StateManager()
    om = OrderManager()
    market_open = False

    # 첫 메시지 처리 (Handshake Hello)
    hello_msg = exchange.read_message()
    if hello_msg and "symbols" in hello_msg:
        for sym_info in hello_msg["symbols"]:
            state.update_position(sym_info["symbol"], sym_info["position"])

    print("봇이 시작되었습니다. P/L 복구 모드 가동.")

    while True:
        message = exchange.read_message()
        if not message: continue

        msg_type = message.get("type")

        if msg_type == "open":
            market_open = True
            print("시장이 열렸습니다.")

        elif msg_type == "book":
            sym = message["symbol"]
            # 최우선 호가 업데이트
            bid = message["buy"][0][0] if message["buy"] else None
            ask = message["sell"][0][0] if message["sell"] else None
            state.update_bid_ask_price(sym, bid, ask)

            # BOND 로직: 1000원 고정 가치를 활용한 안전 수익
            if sym == "BOND" and market_open:
                pos = state.get_position("BOND")
                if pos < 100:
                    exchange.send_add_message(om.next_order(), "BOND", Dir.BUY, 999, 100 - pos)
                if pos > -100:
                    exchange.send_add_message(om.next_order(), "BOND", Dir.SELL, 1001, 100 + pos)

        elif msg_type == "fill":
            # 체결 시 포지션 동기화
            sym, qty, direction = message["symbol"], message["size"], message["dir"]
            delta = qty if direction == Dir.BUY else -qty
            state.update_position(sym, delta)
            print(f"[FILL] {sym} {direction} {qty} | 포지션: {state.positions.get(sym)}")

        elif msg_type == "close":
            market_open = False
            print("시장이 닫혔습니다.")

        elif msg_type == "error":
            print(f"서버 에러: {message.get('error')}")

def parse_arguments():
    parser = argparse.ArgumentParser()
    parser.add_argument("--test", type=str, default="prod-like", choices=["prod-like", "slower", "empty"])
    args = parser.parse_args()
    # 팀 이름에 맞게 호스트네임 설정
    args.exchange_hostname = f"test-exch-{team_name.lower()}"
    args.port = 22000 + (0 if args.test == "prod-like" else 1 if args.test == "slower" else 2)
    return args

if __name__ == "__main__":
    main()