# engine/execution.py
from datetime import datetime, timezone
from indian_paper_trading_strategy.engine.state import load_state, save_state

from indian_paper_trading_strategy.engine.broker import Broker, BrokerAuthExpired
from indian_paper_trading_strategy.engine.ledger import log_event, event_exists
from indian_paper_trading_strategy.engine.db import insert_engine_event, run_with_retry
from indian_paper_trading_strategy.engine.time_utils import compute_logical_time

def _as_float(value):
    if hasattr(value, "item"):
        try:
            return float(value.item())
        except Exception:
            pass
    if hasattr(value, "iloc"):
        try:
            return float(value.iloc[-1])
        except Exception:
            pass
    return float(value)

def _local_tz():
    return datetime.now().astimezone().tzinfo


def _normalize_now(now):
    if now.tzinfo is None:
        return now.replace(tzinfo=_local_tz())
    return now


def _resolve_timing(state, now_ts, sip_interval):
    force_execute = state.get("last_sip_ts") is None
    last = state.get("last_sip_ts") or state.get("last_run")
    if last and not force_execute:
        try:
            last_dt = datetime.fromisoformat(last)
        except ValueError:
            last_dt = None
        if last_dt:
            if last_dt.tzinfo is None:
                last_dt = last_dt.replace(tzinfo=_local_tz())
            if now_ts.tzinfo and last_dt.tzinfo and last_dt.tzinfo != now_ts.tzinfo:
                last_dt = last_dt.astimezone(now_ts.tzinfo)
        if last_dt and (now_ts - last_dt).total_seconds() < sip_interval:
            return False, last, None
    logical_time = compute_logical_time(now_ts, last, sip_interval)
    return True, last, logical_time


def _order_fill(order):
    if not isinstance(order, dict):
        return 0.0, 0.0
    filled_qty = float(order.get("filled_qty") or 0.0)
    average_price = float(order.get("average_price") or order.get("price") or 0.0)
    return filled_qty, average_price


def _apply_filled_orders_to_state(state, orders):
    nifty_filled = 0.0
    gold_filled = 0.0
    total_spent = 0.0

    for order in orders:
        filled_qty, average_price = _order_fill(order)
        if filled_qty <= 0:
            continue
        symbol = (order.get("symbol") or "").upper()
        if symbol.startswith("NIFTYBEES"):
            nifty_filled += filled_qty
        elif symbol.startswith("GOLDBEES"):
            gold_filled += filled_qty
        total_spent += filled_qty * average_price

    if nifty_filled:
        state["nifty_units"] += nifty_filled
    if gold_filled:
        state["gold_units"] += gold_filled
    if total_spent:
        state["total_invested"] += total_spent

    return {
        "nifty_units": nifty_filled,
        "gold_units": gold_filled,
        "amount": total_spent,
    }


def _record_live_order_events(cur, orders, event_ts):
    for order in orders:
        insert_engine_event(cur, "ORDER_PLACED", data=order, ts=event_ts)
        filled_qty, _average_price = _order_fill(order)
        status = (order.get("status") or "").upper()
        if filled_qty > 0:
            insert_engine_event(
                cur,
                "TRADE_EXECUTED",
                data={
                    "order_id": order.get("id"),
                    "symbol": order.get("symbol"),
                    "side": order.get("side"),
                    "qty": filled_qty,
                    "price": order.get("average_price") or order.get("price"),
                },
                ts=event_ts,
            )
            insert_engine_event(cur, "ORDER_FILLED", data={"order_id": order.get("id")}, ts=event_ts)
        elif status == "REJECTED":
            insert_engine_event(cur, "ORDER_REJECTED", data=order, ts=event_ts)
        elif status == "CANCELLED":
            insert_engine_event(cur, "ORDER_CANCELLED", data=order, ts=event_ts)
        elif status == "PENDING":
            insert_engine_event(cur, "ORDER_PENDING", data=order, ts=event_ts)


def _try_execute_sip_paper(
    now,
    market_open,
    sip_interval,
    sip_amount,
    sp_price,
    gd_price,
    eq_w,
    gd_w,
    broker: Broker | None,
    mode: str | None,
):
    def _op(cur, _conn):
        now_ts = _normalize_now(now)
        event_ts = now_ts
        log_event("DEBUG_ENTER_TRY_EXECUTE", {"now": now_ts.isoformat()}, cur=cur, ts=event_ts)

        state = load_state(mode=mode, cur=cur, for_update=True)

        if not market_open:
            return state, False

        should_run, _last, logical_time = _resolve_timing(state, now_ts, sip_interval)
        if not should_run:
            return state, False
        if event_exists("SIP_EXECUTED", logical_time, cur=cur):
            return state, False

        sp_price_val = _as_float(sp_price)
        gd_price_val = _as_float(gd_price)
        eq_w_val = _as_float(eq_w)
        gd_w_val = _as_float(gd_w)
        sip_amount_val = _as_float(sip_amount)

        nifty_qty = (sip_amount_val * eq_w_val) / sp_price_val
        gold_qty = (sip_amount_val * gd_w_val) / gd_price_val

        if broker is None:
            return state, False

        funds = broker.get_funds(cur=cur)
        cash = funds.get("cash")
        if cash is not None and float(cash) < sip_amount_val:
            return state, False

        log_event(
            "DEBUG_EXECUTION_DECISION",
            {
                "last_sip_ts": state.get("last_sip_ts"),
                "now": now_ts.isoformat(),
            },
            cur=cur,
            ts=event_ts,
        )

        orders = [
            broker.place_order(
                "NIFTYBEES.NS",
                "BUY",
                nifty_qty,
                sp_price_val,
                cur=cur,
                logical_time=logical_time,
            ),
            broker.place_order(
                "GOLDBEES.NS",
                "BUY",
                gold_qty,
                gd_price_val,
                cur=cur,
                logical_time=logical_time,
            ),
        ]

        applied = _apply_filled_orders_to_state(state, orders)
        executed = applied["amount"] > 0
        if not executed:
            return state, False

        funds_after = broker.get_funds(cur=cur)
        cash_after = funds_after.get("cash")
        if cash_after is not None:
            state["cash"] = float(cash_after)

        state["last_sip_ts"] = now_ts.isoformat()
        state["last_run"] = now_ts.isoformat()

        save_state(
            state,
            mode=mode,
            cur=cur,
            emit_event=True,
            event_meta={"source": "sip"},
        )

        log_event(
            "SIP_EXECUTED",
            {
                "nifty_units": applied["nifty_units"],
                "gold_units": applied["gold_units"],
                "nifty_price": sp_price_val,
                "gold_price": gd_price_val,
                "amount": applied["amount"],
            },
            cur=cur,
            ts=event_ts,
            logical_time=logical_time,
        )

        return state, True

    return run_with_retry(_op)


def _prepare_live_execution(now_ts, sip_interval, sip_amount_val, sp_price_val, gd_price_val, nifty_qty, gold_qty, mode):
    def _op(cur, _conn):
        state = load_state(mode=mode, cur=cur, for_update=True)
        should_run, _last, logical_time = _resolve_timing(state, now_ts, sip_interval)
        if not should_run:
            return {"ready": False, "state": state}
        if event_exists("SIP_EXECUTED", logical_time, cur=cur):
            return {"ready": False, "state": state}
        if event_exists("SIP_ORDER_ATTEMPTED", logical_time, cur=cur):
            return {"ready": False, "state": state}

        log_event(
            "SIP_ORDER_ATTEMPTED",
            {
                "nifty_units": nifty_qty,
                "gold_units": gold_qty,
                "nifty_price": sp_price_val,
                "gold_price": gd_price_val,
                "amount": sip_amount_val,
            },
            cur=cur,
            ts=now_ts,
            logical_time=logical_time,
        )
        return {"ready": True, "state": state, "logical_time": logical_time}

    return run_with_retry(_op)


def _finalize_live_execution(
    *,
    now_ts,
    mode,
    logical_time,
    orders,
    funds_after,
    sp_price_val,
    gd_price_val,
    auth_failed: bool = False,
    failure_reason: str | None = None,
):
    def _op(cur, _conn):
        state = load_state(mode=mode, cur=cur, for_update=True)
        _record_live_order_events(cur, orders, now_ts)

        applied = _apply_filled_orders_to_state(state, orders)
        executed = applied["amount"] > 0

        if funds_after is not None:
            cash_after = funds_after.get("cash")
            if cash_after is not None:
                state["cash"] = float(cash_after)

        if executed:
            state["last_run"] = now_ts.isoformat()
            state["last_sip_ts"] = now_ts.isoformat()

        save_state(
            state,
            mode=mode,
            cur=cur,
            emit_event=True,
            event_meta={"source": "sip_live"},
        )

        if executed:
            log_event(
                "SIP_EXECUTED",
                {
                    "nifty_units": applied["nifty_units"],
                    "gold_units": applied["gold_units"],
                    "nifty_price": sp_price_val,
                    "gold_price": gd_price_val,
                    "amount": applied["amount"],
                },
                cur=cur,
                ts=now_ts,
                logical_time=logical_time,
            )
        else:
            insert_engine_event(
                cur,
                "SIP_NO_FILL",
                data={
                    "reason": failure_reason or ("broker_auth_expired" if auth_failed else "no_fill"),
                    "orders": orders,
                },
                ts=now_ts,
            )

        return state, executed

    return run_with_retry(_op)


def _try_execute_sip_live(
    now,
    market_open,
    sip_interval,
    sip_amount,
    sp_price,
    gd_price,
    eq_w,
    gd_w,
    broker: Broker | None,
    mode: str | None,
):
    now_ts = _normalize_now(now)
    if not market_open or broker is None:
        return load_state(mode=mode), False

    sp_price_val = _as_float(sp_price)
    gd_price_val = _as_float(gd_price)
    eq_w_val = _as_float(eq_w)
    gd_w_val = _as_float(gd_w)
    sip_amount_val = _as_float(sip_amount)

    nifty_qty = (sip_amount_val * eq_w_val) / sp_price_val
    gold_qty = (sip_amount_val * gd_w_val) / gd_price_val

    prepared = _prepare_live_execution(
        now_ts,
        sip_interval,
        sip_amount_val,
        sp_price_val,
        gd_price_val,
        nifty_qty,
        gold_qty,
        mode,
    )
    if not prepared.get("ready"):
        return prepared.get("state") or load_state(mode=mode), False

    logical_time = prepared["logical_time"]
    orders = []
    funds_after = None
    failure_reason = None
    auth_failed = False

    try:
        funds_before = broker.get_funds()
        cash = funds_before.get("cash")
        if cash is not None and float(cash) < sip_amount_val:
            failure_reason = "insufficient_funds"
        else:
            if nifty_qty > 0:
                orders.append(
                    broker.place_order(
                        "NIFTYBEES.NS",
                        "BUY",
                        nifty_qty,
                        sp_price_val,
                        logical_time=logical_time,
                    )
                )
            if gold_qty > 0:
                orders.append(
                    broker.place_order(
                        "GOLDBEES.NS",
                        "BUY",
                        gold_qty,
                        gd_price_val,
                        logical_time=logical_time,
                    )
                )
        funds_after = broker.get_funds()
    except BrokerAuthExpired:
        auth_failed = True
        try:
            funds_after = broker.get_funds()
        except Exception:
            funds_after = None
    except Exception as exc:
        failure_reason = str(exc)
        try:
            funds_after = broker.get_funds()
        except Exception:
            funds_after = None
        state, _executed = _finalize_live_execution(
            now_ts=now_ts,
            mode=mode,
            logical_time=logical_time,
            orders=orders,
            funds_after=funds_after,
            sp_price_val=sp_price_val,
            gd_price_val=gd_price_val,
            auth_failed=False,
            failure_reason=failure_reason,
        )
        raise

    state, executed = _finalize_live_execution(
        now_ts=now_ts,
        mode=mode,
        logical_time=logical_time,
        orders=orders,
        funds_after=funds_after,
        sp_price_val=sp_price_val,
        gd_price_val=gd_price_val,
        auth_failed=auth_failed,
        failure_reason=failure_reason,
    )
    if auth_failed:
        raise BrokerAuthExpired("Broker session expired during live order execution")
    return state, executed

def try_execute_sip(
    now,
    market_open,
    sip_interval,
    sip_amount,
    sp_price,
    gd_price,
    eq_w,
    gd_w,
    broker: Broker | None = None,
    mode: str | None = "LIVE",
):
    if broker is None:
        return load_state(mode=mode), False
    if getattr(broker, "external_orders", False):
        return _try_execute_sip_live(
            now,
            market_open,
            sip_interval,
            sip_amount,
            sp_price,
            gd_price,
            eq_w,
            gd_w,
            broker,
            mode,
        )
    return _try_execute_sip_paper(
        now,
        market_open,
        sip_interval,
        sip_amount,
        sp_price,
        gd_price,
        eq_w,
        gd_w,
        broker,
        mode,
    )