#!/usr/bin/env python3

import argparse
from enum import Enum
import socket
import json
from state import StateManager

team_name = "HanyangFloorFunction"

# ==================== 설정 ====================
ORDER_SIZE = 2
MAX_POS = 20
ARB_THRESHOLD = 40   

# ====================
class Dir(str, Enum):
    BUY = "BUY"
    SELL = "SELL"

# ====================
def main():
    args = parse_arguments()
    exchange = ExchangeConnection(args)

    state = StateManager()
    hello = exchange.read_message()

    # 포지션 직접 관리
    positions = {}
    for sym in hello["symbols"]:
        positions[sym["symbol"]] = sym["position"]

    order_id = 0
    def next_id():
        nonlocal order_id
        order_id += 1
        return order_id

    active_orders = {}

    # ==================== MARKET MAKING ====================
    def market_make(sym):
        bid = state.bid_prices.get(sym)
        ask = state.ask_prices.get(sym)

        if bid is None or ask is None:
            return

        if ask - bid <= 2: 
            return

        pos = positions.get(sym, 0)

        if abs(pos) > MAX_POS:
            return

        buy_price = bid + 1
        sell_price = ask - 1

        if buy_price >= sell_price:
            return

        # 포지션 조절
        if pos > 0:
            buy_size = 1
            sell_size = ORDER_SIZE + 1
        elif pos < 0:
            buy_size = ORDER_SIZE + 1
            sell_size = 1
        else:
            buy_size = sell_size = ORDER_SIZE

        oid = next_id()
        exchange.send_add_message(oid, sym, Dir.BUY, buy_price, buy_size)
        active_orders[oid] = sym

        oid = next_id()
        exchange.send_add_message(oid, sym, Dir.SELL, sell_price, sell_size)
        active_orders[oid] = sym

    # ==================== BOND 안전 수익 ====================
    def bond_arb():
        bid = state.bid_prices.get("BOND")
        ask = state.ask_prices.get("BOND")

        if ask and ask < 999:
            exchange.send_add_message_ioc(next_id(), "BOND", Dir.BUY, ask, 3)

        if bid and bid > 1001:
            exchange.send_add_message_ioc(next_id(), "BOND", Dir.SELL, bid, 3)

    # ==================== XLF 차익거래 (초안전 버전) ====================
    def xlf_arb():
        bond_ask = state.ask_prices.get("BOND")
        gs_ask   = state.ask_prices.get("GS")
        ms_ask   = state.ask_prices.get("MS")
        wfc_ask  = state.ask_prices.get("WFC")
        xlf_bid  = state.bid_prices.get("XLF")

        bond_bid = state.bid_prices.get("BOND")
        gs_bid   = state.bid_prices.get("GS")
        ms_bid   = state.bid_prices.get("MS")
        wfc_bid  = state.bid_prices.get("WFC")
        xlf_ask  = state.ask_prices.get("XLF")

        if None in [bond_ask, gs_ask, ms_ask, wfc_ask, xlf_bid,
                    bond_bid, gs_bid, ms_bid, wfc_bid, xlf_ask]:
            return

        basket_ask = bond_ask*3 + gs_ask*2 + ms_ask*3 + wfc_ask*2
        basket_bid = bond_bid*3 + gs_bid*2 + ms_bid*3 + wfc_bid*2

        pos = positions.get("XLF", 0)

        if abs(pos) > 10:
            return

        profit1 = xlf_bid * 10 - basket_ask
        if profit1 > ARB_THRESHOLD:
            exchange.send_add_message_ioc(next_id(), "XLF", Dir.SELL, xlf_bid, 10)

        profit2 = basket_bid - xlf_ask * 10
        if profit2 > ARB_THRESHOLD:
            exchange.send_add_message_ioc(next_id(), "XLF", Dir.BUY, xlf_ask, 10)

    # ==================== 리스크 관리 ====================
    def risk():
        for sym in ["GS", "MS", "WFC", "XLF"]:
            pos = positions.get(sym, 0)

            if abs(pos) > MAX_POS:
                if pos > 0:
                    exchange.send_add_message_ioc(
                        next_id(), sym, Dir.SELL,
                        state.bid_prices.get(sym, 1), abs(pos)
                    )
                else:
                    exchange.send_add_message_ioc(
                        next_id(), sym, Dir.BUY,
                        state.ask_prices.get(sym, 99999), abs(pos)
                    )

    # ==================== LOOP ====================
    while True:
        msg = exchange.read_message()

        if msg["type"] == "close":
            break

        elif msg["type"] == "book":
            sym = msg["symbol"]

            bid = msg["buy"][0][0] if msg["buy"] else None
            ask = msg["sell"][0][0] if msg["sell"] else None

            state.update_bid_ask_price(sym, bid, ask)

            for oid in list(active_orders.keys()):
                exchange.send_cancel_message(oid)
                del active_orders[oid]

            # 전략 실행
            bond_arb()
            xlf_arb()

            if sym in ["GS", "MS", "WFC"]:
                market_make(sym)

            risk()

        elif msg["type"] == "fill":
            qty = msg["size"]
            sym = msg["symbol"]

            if msg["dir"] == Dir.BUY:
                positions[sym] = positions.get(sym, 0) + qty
            else:
                positions[sym] = positions.get(sym, 0) - qty

# ====================
class ExchangeConnection:
    def __init__(self, args):
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.connect((args.exchange_hostname, args.port))
        self.reader = s.makefile("r", 1)
        self.writer = s
        self._write({"type": "hello", "team": team_name.upper()})

    def read_message(self):
        msg = json.loads(self.reader.readline())
        if "dir" in msg:
            msg["dir"] = Dir(msg["dir"])
        return msg

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

    def send_add_message_ioc(self, oid, sym, dir, price, size):
        self._write({"type":"add","order_id":oid,"symbol":sym,"dir":dir,"price":price,"size":size,"tif":"IOC"})

    def send_cancel_message(self, oid):
        self._write({"type": "cancel", "order_id": oid})

    def _write(self, msg):
        self.writer.send((json.dumps(msg)+"\n").encode())

# ====================
def parse_arguments():
    parser = argparse.ArgumentParser()

    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("--production", action="store_true")
    group.add_argument("--test", type=str, default="prod-like")

    args = parser.parse_args()

    if args.production:
        args.exchange_hostname = "production"
        args.port = 25000
    else:
        args.exchange_hostname = "test-exch-" + team_name
        args.port = 22000

    return args

if __name__ == "__main__":
    main()