Stock-trading-programming/app/execution/order_executor.py

"""
execution/order_executor.py
Order execution and trade persistence.

DRY_RUN=true means KISClient simulates order fills with the current quote.
"""
import os
import logging
from datetime import datetime

from app.config import FEE_RATE, TAX_RATE
from app.db.models import get_conn
from app.execution.kis_client import KISClient

logger = logging.getLogger(__name__)


class OrderExecutor:
    def __init__(self, kis: KISClient):
        self.kis = kis
        self.dry_run = os.getenv("DRY_RUN", "true").lower() == "true"

    def _calc_fee(self, price: float, qty: int, is_buy: bool) -> float:
        amt = price * qty
        return amt * (FEE_RATE + (0 if is_buy else TAX_RATE))

    async def buy(
        self,
        ticker: str,
        name: str,
        qty: int,
        reason: str = "",
        ai_boosted: bool = False,
        fill_price: float | None = None,
    ) -> dict:
        """Submit a market buy and save the opened trade."""
        try:
            if self.dry_run and fill_price:
                result = {"entry_price": fill_price}
            else:
                result = await self.kis.order_buy(ticker, qty)
            price = result.get("entry_price", 0)
            if not price:
                price = (await self.kis.get_price(ticker))["current"]

            fee = self._calc_fee(price, qty, True)
            trade_id = self._save_trade(
                ticker=ticker,
                name=name,
                entry_price=price,
                qty=qty,
                side="BUY",
                fee=fee,
                ai_boosted=ai_boosted,
            )

            mode = "[DRY]" if self.dry_run else ""
            logger.info("%s BUY %s(%s) %s @ %s", mode, name, ticker, qty, price)
            return {"success": True, "price": price, "qty": qty, "trade_id": trade_id}

        except Exception as e:
            logger.error("BUY failed %s: %s", ticker, e)
            return {"success": False, "error": str(e)}

    async def sell(
        self,
        ticker: str,
        name: str,
        qty: int,
        reason: str = "",
        fill_price: float | None = None,
    ) -> dict:
        """Submit a market sell and save full or partial exit results."""
        try:
            if self.dry_run and fill_price:
                result = {"exit_price": fill_price}
            else:
                result = await self.kis.order_sell(ticker, qty)
            price = result.get("exit_price", 0)
            if not price:
                price = (await self.kis.get_price(ticker))["current"]

            fee = self._calc_fee(price, qty, False)
            self._update_trade_exit(
                ticker=ticker,
                exit_price=price,
                qty=qty,
                reason=reason,
                fee=fee,
            )

            mode = "[DRY]" if self.dry_run else ""
            logger.info("%s SELL %s(%s) %s @ %s [%s]", mode, name, ticker, qty, price, reason)
            return {"success": True, "price": price, "qty": qty}

        except Exception as e:
            logger.error("SELL failed %s: %s", ticker, e)
            return {"success": False, "error": str(e)}

    def _save_trade(self, ticker, name, entry_price, qty, side, fee, ai_boosted=False):
        with get_conn() as conn:
            cur = conn.execute("""
                INSERT INTO trades
                (date, ticker, name, entry_time, entry_price,
                 quantity, side, fee, ai_boosted)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
            """, (
                datetime.now().strftime("%Y-%m-%d"),
                ticker,
                name,
                datetime.now().strftime("%H:%M:%S"),
                entry_price,
                qty,
                side,
                fee,
                1 if ai_boosted else 0,
            ))
            return cur.lastrowid

    def _update_trade_exit(self, ticker, exit_price, qty, reason, fee):
        with get_conn() as conn:
            row = conn.execute("""
                SELECT id, date, name, entry_time, entry_price, quantity,
                       side, fee, ai_boosted
                FROM trades
                WHERE ticker=? AND exit_time IS NULL
                ORDER BY id DESC LIMIT 1
            """, (ticker,)).fetchone()
            if not row:
                logger.warning("No open trade row found for exit: %s", ticker)
                return

            (trade_id, trade_date, name, entry_time, entry_price, trade_qty,
             side, entry_fee, ai_boosted) = row
            actual_qty = qty if qty else trade_qty
            actual_qty = min(actual_qty, trade_qty)
            entry_fee = entry_fee or 0
            allocated_entry_fee = entry_fee * (actual_qty / trade_qty) if trade_qty else 0
            total_fee = allocated_entry_fee + fee
            pnl = (exit_price - entry_price) * actual_qty - total_fee
            exit_time = datetime.now().strftime("%H:%M:%S")

            if actual_qty < trade_qty:
                remaining_qty = trade_qty - actual_qty
                remaining_fee = entry_fee - allocated_entry_fee
                conn.execute("""
                    UPDATE trades
                    SET quantity=?, fee=?
                    WHERE id=?
                """, (remaining_qty, remaining_fee, trade_id))
                conn.execute("""
                    INSERT INTO trades
                    (date, ticker, name, entry_time, exit_time, entry_price,
                     exit_price, quantity, side, exit_reason, pnl, fee,
                     ai_boosted)
                    VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """, (
                    trade_date,
                    ticker,
                    name,
                    entry_time,
                    exit_time,
                    entry_price,
                    exit_price,
                    actual_qty,
                    side,
                    reason,
                    pnl,
                    total_fee,
                    ai_boosted,
                ))
                return

            conn.execute("""
                UPDATE trades
                SET exit_time=?, exit_price=?, exit_reason=?, fee=?, pnl=?
                WHERE id=?
            """, (
                exit_time,
                exit_price,
                reason,
                total_fee,
                pnl,
                trade_id,
            ))