#!/usr/bin/env python3

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

team_name = "HanyangFloorFunction"

# ==================== 설정 ====================
STOCK_SPREAD = 5
ORDER_SIZE = 5
MIN_PROFIT = 5
MAX_POS = 50

# ====================
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
    order_id = 0
    def next_id():
        nonlocal order_id
        order_id += 1
        return order_id

    # ==================== 전략 ====================

    def bond_arb():
        bid = state.bid_prices.get("BOND")
        ask = state.ask_prices.get("BOND")

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

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

    def vale_arb():
        vale_bid = state.bid_prices.get("VALE")
        valbz_ask = state.ask_prices.get("VALBZ")

        if vale_bid and valbz_ask:
            if vale_bid - valbz_ask > MIN_PROFIT:
                exchange.send_add_message_ioc(next_id(), "VALBZ", Dir.BUY, valbz_ask, 1)

    def xlf_arb():
        fair = state.get_fair_value("XLF")
        bid = state.bid_prices.get("XLF")

        if fair and bid:
            if bid - fair > MIN_PROFIT:
                exchange.send_add_message_ioc(next_id(), "XLF", Dir.SELL, bid, 10)

    def market_make(sym):
        mid = state.get_mid_price(sym)
        if mid is None:
            return

        pos = positions.get(sym, 0)

        buy_price = mid - STOCK_SPREAD
        sell_price = mid + STOCK_SPREAD

        if pos > 0:
            buy_size = 1
            sell_size = ORDER_SIZE + 3
        elif pos < 0:
            buy_size = ORDER_SIZE + 3
            sell_size = 1
        else:
            buy_size = sell_size = ORDER_SIZE

        exchange.send_add_message(next_id(), sym, Dir.BUY, buy_price, buy_size)
        exchange.send_add_message(next_id(), sym, Dir.SELL, sell_price, sell_size)

    def risk():
        for sym in ["GS", "MS", "WFC"]:
            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), abs(pos)
                    )
                else:
                    exchange.send_add_message_ioc(
                        next_id(), sym, Dir.BUY,
                        state.ask_prices.get(sym), 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)
            state.update_predicted_price(sym)
            state.update_fair_value()

            bond_arb()
            vale_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 _write(self, msg):
        self.writer.send((json.dumps(msg)+"\n").encode())

# ====================
def parse_arguments():
    parser = argparse.ArgumentParser()
    parser.add_argument("--test", default="prod-like")
    args = parser.parse_args()

    args.exchange_hostname = "test-exch-" + team_name
    args.port = 22000
    return args

if __name__ == "__main__":
    main()