#!/usr/bin/env python3

import argparse
from collections import deque
from enum import Enum
import time
import socket
import json
from state import StateManager
from order import OrderManager

team_name = "HanyangFloorFunction"

MAX_ARBS = 5
XLF_SIZE = 50
VALE_SIZE = 50


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

    hello_message = exchange.read_message()
    print("First message:", hello_message)

    state = StateManager()
    om = OrderManager(exchange)

    market_open = False
    xlf_pending = {}

    def next_id():
        return om.next_order()

    # ---------------- BOND ----------------
    def place_bond_orders():
        if not market_open:
            return

        bid = state.bid_prices["BOND"]
        ask = state.ask_prices["BOND"]

        if bid is None or ask is None:
            return

        buy_price = bid + 1
        sell_price = ask - 1

        size = 100

        exchange.send_add_message(next_id(), "BOND", Dir.BUY, buy_price, size)
        exchange.send_add_message(next_id(), "BOND", Dir.SELL, sell_price, size)

    # ---------------- XLF ----------------
    def try_xlf_arb():
        if not market_open:
            return

        if len(xlf_pending) >= MAX_ARBS:
            return

        bond_ask = state.ask_prices["BOND"]
        gs_ask = state.ask_prices["GS"]
        ms_ask = state.ask_prices["MS"]
        wfc_ask = state.ask_prices["WFC"]
        xlf_bid = state.bid_prices["XLF"]

        bond_bid = state.bid_prices["BOND"]
        gs_bid = state.bid_prices["GS"]
        ms_bid = state.bid_prices["MS"]
        wfc_bid = state.bid_prices["WFC"]
        xlf_ask = state.ask_prices["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

        profit1 = xlf_bid * 10 - basket_ask - 100
        profit2 = basket_bid - xlf_ask * 10 - 100

        # 🔥 공격 (손해도 감수)
        if profit1 > -50:
            for sym, qty in [("BOND",3),("GS",2),("MS",3),("WFC",2)]:
                oid = next_id()
                exchange.send_add_message(oid, sym, Dir.BUY, state.ask_prices[sym], qty)
                xlf_pending[oid] = qty

        elif profit2 > -50:
            oid = next_id()
            exchange.send_add_message(oid, "XLF", Dir.BUY, xlf_ask, XLF_SIZE)
            xlf_pending[oid] = XLF_SIZE

    # ---------------- VALE ----------------
    def try_vale_arb():
        if not market_open:
            return

        vale_bid = state.bid_prices["VALE"]
        vale_ask = state.ask_prices["VALE"]
        valbz_bid = state.bid_prices["VALBZ"]
        valbz_ask = state.ask_prices["VALBZ"]

        if None in [vale_bid, vale_ask, valbz_bid, valbz_ask]:
            return

        profit1 = vale_bid - valbz_ask - 10
        profit2 = valbz_bid - vale_ask - 10

        if profit1 > -5:
            exchange.send_add_message(next_id(), "VALBZ", Dir.BUY, valbz_ask, VALE_SIZE)

        elif profit2 > -5:
            exchange.send_add_message(next_id(), "VALE", Dir.BUY, vale_ask, VALE_SIZE)

    # ---------------- MAIN LOOP ----------------
    while True:
        message = exchange.read_message()

        if message["type"] == "close":
            print("Round ended")
            break

        elif message["type"] == "open":
            print("Market opened")
            market_open = True
            place_bond_orders()

        elif message["type"] == "fill":
            sym = message["symbol"]
            qty = message["size"]
            dir_ = message["dir"]

            if dir_ == Dir.BUY:
                om.update_position(sym, message["order_id"], qty)
            else:
                om.update_position(sym, message["order_id"], -qty)

            print("Position:", om.positions)

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

            state.update_bid_ask_price(
                sym,
                message["buy"][0][0] if message["buy"] else None,
                message["sell"][0][0] if message["sell"] else None
            )

            # 핵심 실행
            try_xlf_arb()
            try_vale_arb()
            place_bond_orders()


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


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

    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, order_id, symbol, dir, price, size):
        self._write({
            "type": "add",
            "order_id": order_id,
            "symbol": symbol,
            "dir": dir,
            "price": price,
            "size": size,
            "tif": "DAY"
        })

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


def parse_arguments():
    parser = argparse.ArgumentParser()
    parser.add_argument("--test", required=True)
    args = parser.parse_args()
    args.exchange_hostname = "test-exch-" + team_name
    args.port = 22000
    return args


if __name__ == "__main__":
    main()