"""
=========================================
CV6 AI Trading OS — Trading State Manager (SINGLE SOURCE OF TRUTH)
=========================================
The ONLY module allowed to own and mutate trading state. Every material
change flows in as ONE TradingEvent via apply_event(); every downstream
consumer (dashboard, API, websocket, trade journal, analytics, risk
sync) READS from here and never keeps its own competing copy.

Design guarantees:
  • Single writer   — apply_event() is the sole mutation entry point,
                      guarded by one re-entrant lock (thread-safe across
                      the FastAPI threadpool + the autonomous engine's
                      background thread).
  • Persistence     — every change is written to SQLite (tsm_orders,
                      tsm_positions, tsm_capital, tsm_events) so full
                      state survives a restart.
  • Recovery        — bootstrap() rebuilds in-memory state from the DB,
                      then reconcile_with_broker() aligns it with the
                      broker's real position book.
  • One-symbol-one-position — positions are keyed by symbol; a same-side
                      fill on an open symbol averages in, an opposite
                      fill reduces/closes. No duplicate positions.
  • Trade journal    — every close writes ONE row to trade_history, so
                      the Journal/PnL/Dashboard read paths finally see
                      real autonomous trades (previously never written).
  • Change feed      — a monotonically increasing version + subscriber
                      hooks let the websocket/dashboard layers push
                      updates without polling broker APIs.
"""

import json
import threading
import time
from typing import Any, Callable, Dict, List, Optional

from loguru import logger

from app.database.session import SessionLocal
from app.state.trading_events import TradingEvent, EventType
from app.state.state_models import (
    OrderRecord, PositionRecord, CapitalState, TradingEventRecord,
)


class TradingStateManager:

    def __init__(self):
        self._lock = threading.RLock()
        # In-memory authoritative state (mirrored to DB on every change)
        self._orders: Dict[str, dict] = {}          # order_id -> order dict
        self._positions: Dict[str, dict] = {}       # symbol   -> OPEN position dict
        self._total_capital     = 0.0
        self._available_capital = 0.0
        self._deployed_capital  = 0.0
        self._realized_pnl      = 0.0
        self._version           = 0                 # bumps on every state change
        self._initialized       = False
        self._subscribers: List[Callable[[dict], None]] = []

    # ── Subscriptions (websocket / dashboard push) ──────────────────────────
    def subscribe(self, callback: Callable[[dict], None]) -> None:
        with self._lock:
            self._subscribers.append(callback)

    def _notify(self, event: TradingEvent) -> None:
        payload = {"event": event.to_dict(), "version": self._version}
        for cb in list(self._subscribers):
            try:
                cb(payload)
            except Exception as e:
                logger.debug(f"[TSM] subscriber error: {e}")

    # ── Bootstrap / recovery ────────────────────────────────────────────────
    def bootstrap(self, default_total_capital: Optional[float] = None) -> dict:
        """
        Rebuild state from the DB after a (re)start. Loads OPEN positions,
        outstanding orders, and the capital snapshot. Idempotent.
        """
        with self._lock:
            db = SessionLocal()
            try:
                # Capital (singleton row id=1)
                cap = db.query(CapitalState).filter(CapitalState.id == 1).first()
                if cap is None:
                    seed = default_total_capital
                    if seed is None:
                        try:
                            from app.autonomous.autonomous_config import load_config
                            seed = load_config().total_capital
                        except Exception:
                            seed = 0.0
                    cap = CapitalState(
                        id=1, total_capital=seed or 0.0,
                        available_capital=seed or 0.0,
                        deployed_capital=0.0, realized_pnl=0.0,
                        updated_at=time.time(),
                    )
                    db.add(cap)
                    db.commit()
                self._total_capital     = cap.total_capital
                self._available_capital = cap.available_capital
                self._deployed_capital  = cap.deployed_capital
                self._realized_pnl      = cap.realized_pnl

                # Open positions
                self._positions.clear()
                for p in db.query(PositionRecord).filter(PositionRecord.status == "OPEN").all():
                    self._positions[p.symbol] = self._pos_row_to_dict(p)

                # Outstanding (non-final) orders
                self._orders.clear()
                for o in db.query(OrderRecord).filter(
                    OrderRecord.status.in_(["PLACED", "PARTIALLY_FILLED"])
                ).all():
                    self._orders[o.order_id] = self._order_row_to_dict(o)

                self._initialized = True
                self._version += 1
                logger.info(
                    f"[TSM] bootstrap: {len(self._positions)} open positions, "
                    f"capital total=₹{self._total_capital:,.0f} "
                    f"available=₹{self._available_capital:,.0f} "
                    f"deployed=₹{self._deployed_capital:,.0f} realized=₹{self._realized_pnl:,.0f}"
                )
                return {"success": True, "open_positions": len(self._positions)}
            except Exception as e:
                logger.error(f"[TSM] bootstrap failed: {e}")
                return {"success": False, "error": str(e)}
            finally:
                db.close()

    def reconcile_with_broker(self, broker_positions: List[dict],
                              close_missing: bool = False) -> dict:
        """
        Align in-memory/DB state with the broker's real position book.
        `broker_positions` is a normalized list of dicts:
          {symbol, exchange, broker, netqty, avg_price, ltp}

        Restore: a real broker position not tracked locally is opened
        (recovery after a restart where the DB was also empty).

        close_missing (intra-session only): when the broker book was
        fetched successfully, force-close any REAL position we still show
        OPEN that the broker reports flat/absent — this catches out-of-band
        exits (manual close in the broker terminal, broker-side square-off,
        SL hit at the exchange). It is OFF by default so a FAILED broker
        fetch (empty list) can never wipe live positions; the caller passes
        close_missing=True only when it knows the fetch succeeded.
        """
        with self._lock:
            restored, closed = 0, 0
            broker_by_symbol = {}
            for bp in broker_positions:
                sym = str(bp.get("symbol", "")).upper()
                if not sym:
                    continue
                qty = int(bp.get("netqty", 0) or 0)
                broker_by_symbol[sym] = qty
                if qty == 0:
                    continue
                if sym not in self._positions:
                    side = "BUY" if qty > 0 else "SELL"
                    avg  = float(bp.get("avg_price", 0) or 0)
                    ltp  = float(bp.get("ltp", avg) or avg)
                    if avg <= 0:
                        continue
                    ev = TradingEvent(
                        event_type=EventType.ORDER_FILLED, symbol=sym,
                        broker=str(bp.get("broker", "")), exchange=str(bp.get("exchange", "NSE")),
                        side=side, filled_quantity=abs(qty), avg_price=avg,
                        fill_status="FILLED", style=str(bp.get("style", "SWING")),
                        strategy="RECONCILED", source="broker_reconcile", mode="REAL",
                    )
                    self._apply_locked(ev, persist=True, notify=False)
                    if sym in self._positions:
                        self._positions[sym]["ltp"] = ltp
                    restored += 1

            if close_missing:
                for sym in list(self._positions.keys()):
                    pos = self._positions[sym]
                    if pos.get("mode") != "REAL":
                        continue
                    broker_qty = broker_by_symbol.get(sym, 0)
                    if broker_qty == 0:  # broker is flat on a symbol we hold open
                        exit_price = pos.get("ltp") or pos.get("entry_price")
                        pnl = self._realized_for(pos, exit_price, pos["qty"])
                        ev = TradingEvent(
                            event_type=EventType.POSITION_CLOSED, symbol=sym,
                            broker=pos.get("broker", ""), exchange=pos.get("exchange", "NSE"),
                            side=pos.get("side", "BUY"), filled_quantity=pos["qty"],
                            avg_price=exit_price, pnl=pnl, reason="BROKER_RECONCILE_FLAT",
                            style=pos.get("style", ""), strategy=pos.get("strategy", ""),
                            mode="REAL", source="broker_reconcile",
                        )
                        self._apply_locked(ev, persist=True, notify=True)
                        closed += 1

            self._version += 1
            if restored or closed:
                logger.info(f"[TSM] broker reconcile: restored={restored} closed={closed}")
            return {"success": True, "restored": restored, "closed": closed}

    # ── The single mutation entry point ─────────────────────────────────────
    def apply_event(self, event: TradingEvent) -> dict:
        with self._lock:
            result = self._apply_locked(event, persist=True, notify=True)
            return result

    def _apply_locked(self, event: TradingEvent, persist: bool, notify: bool) -> dict:
        et = event.event_type
        if et in (EventType.ORDER_FILLED, EventType.ORDER_PARTIALLY_FILLED):
            res = self._handle_fill(event)
        elif et == EventType.POSITION_CLOSED:
            res = self._handle_close(event)
        elif et == EventType.ORDER_REJECTED:
            res = self._handle_reject(event)
        elif et == EventType.PRICE_TICK:
            res = self._handle_tick(event)
        else:
            res = {"success": False, "error": f"unknown event {et}"}

        if persist and et != EventType.PRICE_TICK:
            self._persist_event(event)
            self._persist_capital()
        if et != EventType.PRICE_TICK:
            self._version += 1
        if notify:
            self._notify(event)
        return res

    # ── Handlers ────────────────────────────────────────────────────────────
    def _handle_fill(self, ev: TradingEvent) -> dict:
        sym = ev.symbol.upper()
        qty = ev.filled_quantity or ev.quantity
        price = ev.avg_price or ev.price
        if qty <= 0 or price <= 0:
            return {"success": False, "error": "fill has zero qty/price"}

        # order record
        self._orders[ev.order_id or f"{sym}-{int(ev.ts)}"] = {
            "order_id": ev.order_id, "symbol": sym, "broker": ev.broker,
            "exchange": ev.exchange, "side": ev.side, "quantity": ev.quantity,
            "filled_quantity": qty, "avg_price": price,
            "status": ev.fill_status or "FILLED", "mode": ev.mode,
            "strategy": ev.strategy, "ai_confidence": ev.ai_confidence,
        }

        existing = self._positions.get(sym)
        value = price * qty
        if existing is None:
            self._positions[sym] = {
                "symbol": sym, "broker": ev.broker, "exchange": ev.exchange,
                "style": ev.style or "INTRADAY", "side": ev.side, "qty": qty,
                "entry_price": price, "ltp": price,
                "stop_loss": ev.stop_loss, "target_price": ev.target_price,
                "status": "OPEN", "strategy": ev.strategy,
                "ai_confidence": ev.ai_confidence, "mode": ev.mode,
                "entry_order_id": ev.order_id, "sl_order_id": ev.sl_order_id,
                "opened_at": ev.ts,
            }
            self._deployed_capital  += value
            self._available_capital -= value
        elif existing["side"] == ev.side:
            # same-side add — weighted average, one position preserved
            old_val = existing["entry_price"] * existing["qty"]
            new_qty = existing["qty"] + qty
            existing["entry_price"] = (old_val + value) / new_qty if new_qty else price
            existing["qty"] = new_qty
            existing["ltp"] = price
            self._deployed_capital  += value
            self._available_capital -= value
        else:
            # opposite-side fill reduces/closes — treat as (partial) close
            close_qty = min(existing["qty"], qty)
            pnl = self._realized_for(existing, price, close_qty)
            self._book_close(existing, price, close_qty, pnl, ev.reason or "OPPOSITE_FILL", ev)
            # any residual opposite qty opens a new position the other way
            residual = qty - close_qty
            if residual > 0:
                self._positions[sym] = {
                    "symbol": sym, "broker": ev.broker, "exchange": ev.exchange,
                    "style": ev.style or "INTRADAY", "side": ev.side, "qty": residual,
                    "entry_price": price, "ltp": price,
                    "stop_loss": ev.stop_loss, "target_price": ev.target_price,
                    "status": "OPEN", "strategy": ev.strategy,
                    "ai_confidence": ev.ai_confidence, "mode": ev.mode,
                    "entry_order_id": ev.order_id, "sl_order_id": ev.sl_order_id,
                    "opened_at": ev.ts,
                }
                rvalue = price * residual
                self._deployed_capital  += rvalue
                self._available_capital -= rvalue

        self._persist_order(ev, qty, price)
        if sym in self._positions:
            self._persist_position(self._positions[sym])
        return {"success": True, "symbol": sym, "qty": qty}

    def _handle_close(self, ev: TradingEvent) -> dict:
        sym = ev.symbol.upper()
        pos = self._positions.get(sym)
        if pos is None:
            return {"success": False, "error": f"no open position for {sym}"}
        exit_price = ev.avg_price or ev.price or pos["ltp"]
        close_qty  = ev.filled_quantity or pos["qty"]
        close_qty  = min(close_qty, pos["qty"])
        pnl = ev.pnl if ev.pnl else self._realized_for(pos, exit_price, close_qty)
        self._book_close(pos, exit_price, close_qty, pnl, ev.reason or "CLOSE", ev)
        return {"success": True, "symbol": sym, "pnl": pnl}

    def _book_close(self, pos: dict, exit_price: float, close_qty: int,
                    pnl: float, reason: str, ev: TradingEvent) -> None:
        sym = pos["symbol"]
        entry_value = pos["entry_price"] * close_qty
        self._realized_pnl      += pnl
        self._available_capital += entry_value + pnl
        self._deployed_capital  -= entry_value
        remaining = pos["qty"] - close_qty
        if remaining > 0:
            pos["qty"] = remaining
            self._persist_position(pos)
        else:
            self._positions.pop(sym, None)
            self._persist_position_closed(pos, exit_price, pnl, reason, ev)
        # Trade journal — write ONE row so /history + Dashboard see it
        self._write_journal(pos, exit_price, close_qty, pnl, reason, ev)

    def _handle_reject(self, ev: TradingEvent) -> dict:
        self._persist_order(ev, 0, ev.price, status="REJECTED")
        return {"success": True, "rejected": ev.symbol}

    def _handle_tick(self, ev: TradingEvent) -> dict:
        pos = self._positions.get(ev.symbol.upper())
        if pos and ev.ltp > 0:
            pos["ltp"] = ev.ltp
        return {"success": True}

    # ── Derived reads ───────────────────────────────────────────────────────
    def _realized_for(self, pos: dict, exit_price: float, qty: int) -> float:
        if pos["side"] == "BUY":
            return (exit_price - pos["entry_price"]) * qty
        return (pos["entry_price"] - exit_price) * qty

    def _unrealized(self) -> float:
        total = 0.0
        for p in self._positions.values():
            if p["side"] == "BUY":
                total += (p["ltp"] - p["entry_price"]) * p["qty"]
            else:
                total += (p["entry_price"] - p["ltp"]) * p["qty"]
        return total

    def snapshot(self) -> dict:
        with self._lock:
            unreal = self._unrealized()
            positions = []
            for p in self._positions.values():
                pnl = ((p["ltp"] - p["entry_price"]) if p["side"] == "BUY"
                       else (p["entry_price"] - p["ltp"])) * p["qty"]
                positions.append({
                    "symbol": p["symbol"], "exchange": p["exchange"],
                    "broker": p["broker"], "style": p["style"], "side": p["side"],
                    "quantity": p["qty"], "avg_price": round(p["entry_price"], 2),
                    "ltp": round(p["ltp"], 2), "pnl": round(pnl, 2),
                    "pnl_pct": round(pnl / (p["entry_price"] * p["qty"]) * 100, 2)
                               if p["entry_price"] and p["qty"] else 0.0,
                    "stop_loss": round(p["stop_loss"], 2), "target": round(p["target_price"], 2),
                    "strategy": p["strategy"], "mode": p["mode"],
                })
            return {
                "source": "trading_state_manager",
                "version": self._version,
                "initialized": self._initialized,
                "positions": positions,
                "position_count": len(positions),
                "orders": list(self._orders.values()),
                "capital": {
                    "total": round(self._total_capital, 2),
                    "available": round(self._available_capital, 2),
                    "deployed": round(self._deployed_capital, 2),
                },
                "pnl": {
                    "realized": round(self._realized_pnl, 2),
                    "unrealized": round(unreal, 2),
                    "total": round(self._realized_pnl + unreal, 2),
                },
            }

    def get_position(self, symbol: str) -> Optional[dict]:
        with self._lock:
            return self._positions.get(symbol.upper())

    def has_open_position(self, symbol: str) -> bool:
        with self._lock:
            return symbol.upper() in self._positions

    def version(self) -> int:
        return self._version

    # ── Persistence helpers (each uses its own short-lived session) ─────────
    def _pos_row_to_dict(self, p: PositionRecord) -> dict:
        return {
            "symbol": p.symbol, "broker": p.broker, "exchange": p.exchange,
            "style": p.style, "side": p.side, "qty": p.qty,
            "entry_price": p.entry_price, "ltp": p.ltp or p.entry_price,
            "stop_loss": p.stop_loss, "target_price": p.target_price,
            "status": p.status, "strategy": p.strategy,
            "ai_confidence": p.ai_confidence, "mode": p.mode,
            "entry_order_id": p.entry_order_id, "sl_order_id": p.sl_order_id or "",
            "opened_at": p.opened_at,
            "_db_id": p.id,
        }

    def _order_row_to_dict(self, o: OrderRecord) -> dict:
        return {
            "order_id": o.order_id, "symbol": o.symbol, "broker": o.broker,
            "exchange": o.exchange, "side": o.side, "quantity": o.quantity,
            "filled_quantity": o.filled_quantity, "avg_price": o.avg_price,
            "status": o.status, "mode": o.mode, "strategy": o.strategy,
            "ai_confidence": o.ai_confidence,
        }

    def _persist_order(self, ev: TradingEvent, filled_qty: int, price: float,
                       status: Optional[str] = None) -> None:
        oid = ev.order_id or f"{ev.symbol.upper()}-{int(ev.ts)}"
        db = SessionLocal()
        try:
            rec = db.query(OrderRecord).filter(OrderRecord.order_id == oid).first()
            st = status or ev.fill_status or "FILLED"
            if rec is None:
                db.add(OrderRecord(
                    order_id=oid, broker=ev.broker, symbol=ev.symbol.upper(),
                    exchange=ev.exchange, side=ev.side, product=ev.style or "INTRADAY",
                    quantity=ev.quantity, filled_quantity=filled_qty, price=ev.price,
                    avg_price=price, status=st, mode=ev.mode, strategy=ev.strategy,
                    ai_confidence=ev.ai_confidence, reason=ev.reason,
                    created_at=ev.ts, updated_at=time.time(),
                ))
            else:
                rec.filled_quantity = filled_qty
                rec.avg_price = price
                rec.status = st
                rec.updated_at = time.time()
            db.commit()
        except Exception as e:
            logger.warning(f"[TSM] persist order failed: {e}")
            db.rollback()
        finally:
            db.close()

    def _persist_position(self, pos: dict) -> None:
        db = SessionLocal()
        try:
            rec = None
            if pos.get("_db_id"):
                rec = db.query(PositionRecord).filter(PositionRecord.id == pos["_db_id"]).first()
            if rec is None:
                rec = db.query(PositionRecord).filter(
                    PositionRecord.symbol == pos["symbol"],
                    PositionRecord.status == "OPEN",
                ).first()
            if rec is None:
                rec = PositionRecord(symbol=pos["symbol"], opened_at=pos.get("opened_at", time.time()))
                db.add(rec)
            rec.broker = pos["broker"]; rec.exchange = pos["exchange"]
            rec.style = pos["style"]; rec.side = pos["side"]; rec.qty = pos["qty"]
            rec.entry_price = pos["entry_price"]; rec.ltp = pos["ltp"]
            rec.stop_loss = pos["stop_loss"]; rec.target_price = pos["target_price"]
            rec.status = "OPEN"; rec.strategy = pos["strategy"]
            rec.ai_confidence = pos["ai_confidence"]; rec.mode = pos["mode"]
            rec.entry_order_id = pos.get("entry_order_id", "")
            rec.sl_order_id = pos.get("sl_order_id", "")
            rec.updated_at = time.time()
            db.commit()
            pos["_db_id"] = rec.id
        except Exception as e:
            logger.warning(f"[TSM] persist position failed: {e}")
            db.rollback()
        finally:
            db.close()

    def _persist_position_closed(self, pos: dict, exit_price: float, pnl: float,
                                 reason: str, ev: TradingEvent) -> None:
        db = SessionLocal()
        try:
            rec = None
            if pos.get("_db_id"):
                rec = db.query(PositionRecord).filter(PositionRecord.id == pos["_db_id"]).first()
            if rec is None:
                rec = db.query(PositionRecord).filter(
                    PositionRecord.symbol == pos["symbol"], PositionRecord.status == "OPEN",
                ).first()
            if rec is None:
                rec = PositionRecord(symbol=pos["symbol"], opened_at=pos.get("opened_at", time.time()),
                                     entry_price=pos["entry_price"], side=pos["side"], qty=pos["qty"])
                db.add(rec)
            rec.status = "CLOSED"; rec.exit_price = exit_price
            rec.realized_pnl = pnl; rec.exit_reason = reason
            rec.closed_at = time.time(); rec.updated_at = time.time()
            db.commit()
        except Exception as e:
            logger.warning(f"[TSM] persist position-close failed: {e}")
            db.rollback()
        finally:
            db.close()

    def _persist_capital(self) -> None:
        db = SessionLocal()
        try:
            cap = db.query(CapitalState).filter(CapitalState.id == 1).first()
            if cap is None:
                cap = CapitalState(id=1)
                db.add(cap)
            cap.total_capital = self._total_capital
            cap.available_capital = self._available_capital
            cap.deployed_capital = self._deployed_capital
            cap.realized_pnl = self._realized_pnl
            cap.updated_at = time.time()
            db.commit()
        except Exception as e:
            logger.warning(f"[TSM] persist capital failed: {e}")
            db.rollback()
        finally:
            db.close()

    def _persist_event(self, ev: TradingEvent) -> None:
        db = SessionLocal()
        try:
            db.add(TradingEventRecord(
                event_id=ev.event_id, event_type=ev.event_type.value,
                symbol=ev.symbol.upper(), broker=ev.broker, order_id=ev.order_id,
                payload=json.dumps(ev.to_dict(), default=str), ts=ev.ts,
            ))
            db.commit()
        except Exception as e:
            logger.warning(f"[TSM] persist event failed: {e}")
            db.rollback()
        finally:
            db.close()

    def _write_journal(self, pos: dict, exit_price: float, qty: int,
                       pnl: float, reason: str, ev: TradingEvent) -> None:
        """Write ONE closed-trade row to trade_history (the Trade Journal)."""
        try:
            from app.history.history_service import history_service
            from app.history.history_schemas import TradeCreate
            db = SessionLocal()
            try:
                history_service.create(db, TradeCreate(
                    broker=pos["broker"], symbol=pos["symbol"], exchange=pos["exchange"],
                    trade_type="SELL" if pos["side"] == "BUY" else "BUY",
                    status="EXECUTED", quantity=qty,
                    price=exit_price, avg_price=exit_price, filled_qty=qty,
                    product="MIS" if pos["style"] in ("INTRADAY", "SCALPING") else "NRML",
                    order_type="MARKET", pnl=round(pnl, 2), charges=ev.charges,
                    ai_confidence=pos.get("ai_confidence", 0.0), strategy=pos.get("strategy", ""),
                    notes=f"{reason}: entry ₹{pos['entry_price']:.2f} exit ₹{exit_price:.2f}",
                    placed_at=time.time(),
                ))
            finally:
                db.close()
        except Exception as e:
            logger.warning(f"[TSM] journal write failed: {e}")


# ── Singleton ────────────────────────────────────────────────────────────────
trading_state_manager = TradingStateManager()