#!/usr/bin/env python3
"""
ORB FLEET WORKER
Drop-in distributed mode for the Retest ORB Signal Bot.

Place this file + fleet_client.py NEXT TO orb_signal_bot.py on every laptop.
It reuses the entire existing pipeline (Chrome thread, OCR worker thread,
validator, price monitor, Telegram) — the only things that change:

  1. Tickers come from the fleet coordinator (atomic lease claims),
     not from the local tradingview_links.txt.
  2. Before any Telegram alert, the coordinator's GLOBAL signal registry
     is consulted → exactly one laptop alerts per signal per day.
  3. A heartbeat thread reports liveness + obeys dashboard commands
     (pause / batch size / re-scan interval).

Run:  python orb_fleet_worker.py
Config: add a "fleet" section to orb_config.json (see below).

    "fleet": {
        "enabled": true,
        "coordinator_url": "https://dms.burraqvehicles.com/orb_fleet/orb_fleet_api.php",
        "api_key": "CHANGE_ME_long_random_worker_key",
        "worker_name": "laptop-1",
        "heartbeat_seconds": 30,
        "idle_wait_seconds": 15,
        "signal_failopen": true
    }

Failure model:
  • This laptop dies      → its leases expire server-side; others take over.
  • Coordinator dies      → worker keeps retrying claims; if signal_failopen
                            is true, any in-flight signal still alerts locally.
  • New laptop starts     → it simply begins claiming; load self-balances.
"""

import os
import sys
import json
import time
import queue
import threading
from datetime import datetime

# Reuse everything from the existing bot — must be in the same folder
from orb_signal_bot import ORBSignalMonitor
from fleet_client   import FleetClient

from io import BytesIO
from PIL import Image
import base64

WORKER_VERSION = "1.0"


class FleetORBWorker(ORBSignalMonitor):

    def __init__(self, config_file: str = "orb_config.json"):
        super().__init__(config_file=config_file, links_file="tradingview_links.txt")

        with open(config_file) as f:
            raw = json.load(f)
        fc = raw.get("fleet", {})
        if not fc.get("enabled"):
            raise RuntimeError(
                'orb_config.json has no enabled "fleet" section — '
                "add it or run the standalone orb_signal_bot.py instead."
            )

        self.fleet = FleetClient(
            base_url    = fc["coordinator_url"],
            api_key     = fc["api_key"],
            worker_name = fc.get("worker_name") or os.environ.get("COMPUTERNAME", "worker"),
            logger      = self.logger,
        )
        self.hb_seconds     = int(fc.get("heartbeat_seconds", 30))
        self.idle_wait      = int(fc.get("idle_wait_seconds", 15))
        self.failopen       = bool(fc.get("signal_failopen", True))

        # Live commands from coordinator (updated by heartbeat thread)
        self._commands      = {"paused": False, "batch_size": 4, "check_interval": 30}
        self._cmd_lock      = threading.Lock()
        self._current_ticker: str = ""
        self._stop          = threading.Event()

        # local signal id → fleet signal id (for marking telegram_sent)
        self._fleet_sig_map = {}

        # Local bad-scan streak (blank charts / errors) — mirrors the server's
        # quarantine counter so we can upload one evidence screenshot early
        self._local_bad_streak = 0

        # Hook GLOBAL dedup into both DB connections (Chrome + OCR threads)
        self._install_signal_hooks(self.db)
        self._install_signal_hooks(self._worker_db)

        # Hook price-monitor outcomes → fleet win-rate table
        self._install_outcome_hook()

        # Hook price polls → coordinator paper-trading engine (throttled)
        self._install_price_push_hook()

        # Hook validator rejections → anomaly screenshot upload
        self._install_validator_hook()

        self.logger.info("🛰  Fleet worker mode initialised")

    # ─────────────────────────────────────────────────────────────────────────
    # Global dedup hooks — wrap save_signal / mark_telegram_sent on the
    # ORBDatabaseManager instances so orb_signal_bot's pipeline needs ZERO edits.
    # ─────────────────────────────────────────────────────────────────────────

    def _install_signal_hooks(self, db_obj):
        orig_save = db_obj.save_signal
        orig_mark = db_obj.mark_telegram_sent
        worker    = self

        def save_signal(ticker, signal_data, screenshot_path=None, trend_snapshot=None):
            new_id = orig_save(ticker, signal_data, screenshot_path,
                               trend_snapshot=trend_snapshot)
            # Local dup / DB error → nothing more to do
            if not isinstance(new_id, int) or new_id <= 0:
                return new_id

            # Locally new → ask the fleet if it's GLOBALLY new
            resp = worker.fleet.signal_check(
                ticker      = ticker,
                signal_type = str(signal_data.get("signal_type", "")).upper(),
                entry       = _flt(signal_data.get("entry")),
                tp          = _flt(signal_data.get("tp")),
                sl          = _flt(signal_data.get("sl")),
                sig_time    = signal_data.get("time"),
            )

            if resp is None:
                # Coordinator unreachable
                if worker.failopen:
                    worker.logger.warning(
                        f"[{ticker}] ⚠ Fleet unreachable for signal_check — "
                        f"FAIL-OPEN: alerting locally (rare dup possible)")
                    return new_id
                worker.logger.warning(
                    f"[{ticker}] ⚠ Fleet unreachable — FAIL-CLOSED: "
                    f"suppressing alert (signal saved locally, id={new_id})")
                return -1

            if resp.get("duplicate"):
                worker.logger.info(
                    f"[{ticker}] 🌐 Fleet duplicate — first seen by "
                    f"'{resp.get('first_seen_by')}' (fleet id={resp.get('signal_id')}) "
                    f"— suppressing Telegram + monitor on this node")
                return -1  # bot treats as duplicate: no Telegram, no price monitor

            worker._fleet_sig_map[new_id] = int(resp.get("signal_id", 0))
            worker.logger.info(
                f"[{ticker}] 🌐 Fleet-NEW signal (fleet id={resp.get('signal_id')}) "
                f"— this node owns the alert")
            return new_id

        def mark_telegram_sent(signal_id):
            ok = orig_mark(signal_id)
            fid = worker._fleet_sig_map.get(signal_id)
            if fid:
                worker.fleet.signal_tg_sent(fid)
            return ok

        db_obj.save_signal        = save_signal
        db_obj.mark_telegram_sent = mark_telegram_sent

    # ─────────────────────────────────────────────────────────────────────────
    # Outcome sync — when a price monitor concludes (tp_hit / sl_hit / timeout)
    # report it to the coordinator so the dashboard's win-rate table is live.
    # Wraps _update_monitor_status because BOTH the TP/SL path and the timeout
    # path go through it with an `outcome` kwarg.
    # ─────────────────────────────────────────────────────────────────────────

    def _install_outcome_hook(self):
        pm     = self.price_monitor
        orig   = pm._update_monitor_status
        worker = self

        def _update_monitor_status(signal_id, status, outcome=None,
                                   max_fav=None, max_fav_pct=None):
            r = orig(signal_id, status, outcome=outcome,
                     max_fav=max_fav, max_fav_pct=max_fav_pct)
            if outcome:
                fid = worker._fleet_sig_map.get(signal_id)
                if fid:
                    try:
                        worker.fleet.signal_outcome(
                            fid, outcome, exit_price=None, max_fav_pct=max_fav_pct)
                        worker.logger.info(
                            f"🛰  Outcome '{outcome}' synced to fleet (signal {fid})")
                    except Exception as _oe:
                        worker.logger.debug(f"Fleet outcome sync failed: {_oe}")
            return r

        pm._update_monitor_status = _update_monitor_status

    # ─────────────────────────────────────────────────────────────────────────
    # Price push — the price monitors already poll live prices for their
    # signals; mirror each poll to the coordinator (throttled per ticker) so
    # the server-side paper-trading engine has fresh data without its own
    # market-data subscription. Fire-and-forget threads, never blocks polls.
    # ─────────────────────────────────────────────────────────────────────────

    def _install_price_push_hook(self):
        pm     = self.price_monitor
        orig   = pm._save_price_to_db
        worker = self
        throttle = {}          # ticker -> last push monotonic ts
        MIN_GAP  = 20.0        # seconds between pushes per ticker

        def _save_price_to_db(ticker, price):
            r = orig(ticker, price)
            try:
                now = time.monotonic()
                if price and now - throttle.get(ticker, 0) >= MIN_GAP:
                    throttle[ticker] = now
                    threading.Thread(
                        target=worker.fleet.price_push, args=(ticker, float(price)),
                        name="PricePush", daemon=True).start()
            except Exception:
                pass
            return r

        pm._save_price_to_db = _save_price_to_db

    # ─────────────────────────────────────────────────────────────────────────
    # Anomaly screenshots — when the validator REJECTS a detected signal
    # (chased / already stopped), upload the chart screenshot so misreads can
    # be reviewed from the dashboard instead of remoting into this laptop.
    # ─────────────────────────────────────────────────────────────────────────

    def _install_validator_hook(self):
        val    = self.validator
        orig   = val.validate
        worker = self

        def validate(signal, chart_url, ticker):
            outcome, live_price, reason = orig(signal, chart_url, ticker)
            try:
                oname = str(outcome).lower()
                if oname in ("chased", "stopped"):
                    ss_path = os.path.join(worker.screenshot_dir, f"{ticker}.png")
                    worker._upload_screenshot_async(
                        ticker, oname, ss_path,
                        note=f"{signal.get('signal_type','?')} entry={signal.get('entry','?')} — {reason}"[:250])
            except Exception:
                pass
            return outcome, live_price, reason

        val.validate = validate

    def _upload_screenshot_async(self, ticker: str, reason: str,
                                 png_path: str, note: str = None):
        """Downscale → JPEG → base64 → upload, off the OCR thread."""
        def _job():
            try:
                if not os.path.exists(png_path):
                    return
                img = Image.open(png_path)
                if img.mode != "RGB":
                    img = img.convert("RGB")
                if img.width > 900:
                    img = img.resize((900, int(img.height * 900 / img.width)))
                buf = BytesIO()
                img.save(buf, format="JPEG", quality=68)
                b64 = base64.b64encode(buf.getvalue()).decode("ascii")
                self.fleet.screenshot_upload(ticker, reason, b64, note=note)
                self.logger.info(f"[{ticker}] 📸 Anomaly screenshot uploaded ({reason})")
            except Exception as e:
                self.logger.debug(f"[{ticker}] screenshot upload failed: {e}")
        threading.Thread(target=_job, name="ScreenUpload", daemon=True).start()

    @staticmethod
    def _looks_blank(img: Image.Image) -> bool:
        """Blind-scanner heuristic: a real TradingView chart is high-contrast
        (candles, gridlines, axis text). A logged-out page, captcha, or failed
        render is mostly uniform → very low grayscale std-dev."""
        try:
            from PIL import ImageStat
            g = img.convert("L")
            # Sample down for speed — std-dev is scale-invariant enough
            g.thumbnail((400, 400))
            return ImageStat.Stat(g).stddev[0] < 10.0
        except Exception:
            return False  # never quarantine on heuristic failure

    # ─────────────────────────────────────────────────────────────────────────
    # Heartbeat thread
    # ─────────────────────────────────────────────────────────────────────────

    def _heartbeat_loop(self):
        while not self._stop.is_set():
            cmd = self.fleet.heartbeat(current_ticker=self._current_ticker or None)
            if cmd:
                with self._cmd_lock:
                    changed = cmd != self._commands
                    self._commands = cmd
                if changed:
                    self.logger.info(f"🛰  Fleet commands updated: {cmd}")
            self._stop.wait(self.hb_seconds)

    def _cmd(self, key, default=None):
        with self._cmd_lock:
            return self._commands.get(key, default)

    # ─────────────────────────────────────────────────────────────────────────
    # Fleet main loop (replaces the static-links _loop)
    # ─────────────────────────────────────────────────────────────────────────

    def start(self) -> bool:
        self.logger.info("=" * 60)
        self.logger.info("ORB FLEET WORKER — distributed scanning")
        self.logger.info("=" * 60)
        self.stats["start_time"] = datetime.now()

        if not self.chrome.connect():
            self.logger.error("Cannot connect to Chrome"); return False
        if not self.db.connect():
            self.logger.error("Cannot connect to local MySQL"); return False
        if not self.trend_db.connect():
            self.logger.warning("trend_info DB connect failed — supertrend data won't save")
        if self.telegram and not self.telegram.test_connection():
            self.logger.warning("Telegram test failed — signals still recorded")

        # Register with the coordinator (retry until reachable)
        while not self.fleet.register(version=WORKER_VERSION):
            self.logger.error("Coordinator unreachable — retrying in 20s "
                              "(check fleet.coordinator_url / api_key)")
            time.sleep(20)

        threading.Thread(target=self._heartbeat_loop,
                         name="FleetHeartbeat", daemon=True).start()

        # Start the existing OCR pipeline thread
        if self._worker_db.connect():
            self._worker_trend_db.connect()
            # Re-hook: connect() may have been called after hooks — hooks are on
            # the instance so they survive, nothing to do.
            self._pipeline_enabled = True
        else:
            self._pipeline_enabled = False
            self.logger.warning("Secondary DB connect failed — sequential OCR mode")

        self._ocr_queue = queue.Queue(maxsize=4)
        threading.Thread(target=self._ocr_worker, name="OCRWorker",
                         daemon=True).start()
        self.logger.info("✅ Pipeline mode active" if self._pipeline_enabled
                         else "⚠ Sequential mode")

        try:
            self._fleet_loop()
        except KeyboardInterrupt:
            self.logger.info("Stopped by user")
        finally:
            self._stop.set()
            try:
                self.fleet.release_all()
                self.logger.info("🛰  Leases released back to fleet")
            except Exception:
                pass
            self._cleanup()
        return True

    def _fleet_loop(self):
        batches = 0
        last_pause_reason = None
        while True:
            if self._cmd("paused"):
                self._current_ticker = ""
                reason = self._cmd("pause_reason") or "paused by dashboard"
                if reason != last_pause_reason:
                    icon = "🌙" if "market closed" in reason else ("🟡" if "QUARANTINED" in reason else "⏸")
                    self.logger.info(f"{icon} Idle — {reason}")
                    last_pause_reason = reason
                    self.fleet.progress(None, "idle")
                # Market-closed naps can be long; stay responsive to heartbeat cmds
                time.sleep(min(60, self.idle_wait * 2) if "market closed" in reason
                           else self.idle_wait)
                continue
            last_pause_reason = None

            if not self.chrome.is_connected():
                self.logger.warning("Chrome disconnected – reconnecting…")
                if not self.chrome.reconnect():
                    time.sleep(30); continue
            if not self.db.is_connected():
                self.logger.warning("Local DB disconnected – reconnecting…")
                if not self.db.reconnect():
                    time.sleep(30); continue

            claim = self.fleet.claim(batch_size=self._cmd("batch_size", 4))
            if claim is None:
                self.logger.warning("Coordinator unreachable — retrying in 20s")
                time.sleep(20); continue

            if claim.get("paused"):
                # Claim-side pause (e.g. market just closed between heartbeats)
                wait = int(claim.get("retry_in") or self.idle_wait)
                reason = claim.get("pause_reason") or "paused"
                if reason != last_pause_reason:
                    self.logger.info(f"⏸ Coordinator says: {reason} — napping {wait}s")
                    last_pause_reason = reason
                time.sleep(wait)
                continue

            tickers = claim.get("tickers", [])
            if not tickers:
                # Pool empty (all leased or all freshly scanned) → drain OCR, nap
                self._drain_ocr_queue()
                wait = int(claim.get("retry_in") or self.idle_wait)
                self.logger.info(f"💤 No claimable tickers — napping {wait}s")
                self._current_ticker = ""
                time.sleep(wait)
                continue

            token = claim.get("lease_token", "")
            batches += 1
            self.stats["cycles"] += 1
            turbo = " 🚀 TURBO" if claim.get("turbo") else ""
            self.logger.info(
                f"\n{'='*60}\n"
                f"Batch #{batches}{turbo} — claimed {len(tickers)}: "
                f"{', '.join(t['ticker'] for t in tickers)}\n"
                f"{'='*60}")

            for i, t in enumerate(tickers, 1):
                if self._cmd("paused"):
                    break
                self._scan_one(i, len(tickers), t, token)

            # Periodically let the OCR queue fully drain so report ordering
            # and Telegram sends never lag far behind Chrome
            if batches % 3 == 0:
                self._drain_ocr_queue()

    def _scan_one(self, idx: int, total: int, t: dict, token: str):
        ticker, link, tid = t["ticker"], t["tv_url"], int(t["id"])
        self._current_ticker = ticker
        self.fleet.progress(ticker, "loading")   # → dashboard "Now scanning"
        self.logger.info(f"[{idx}/{total}] {ticker} — {link[-70:]}")

        t0 = time.time()
        tab_id, err, boxes_found = None, None, False
        ocr_ok = True
        try:
            tab_id = self.chrome.open_new_tab(link)
            if not tab_id:
                err = "tab open failed"
                self.stats["errors"] += 1
                return

            tab_info = {"id": tab_id, "symbol": ticker}
            self.chrome.wait_for_page_ready(tab_info)
            time.sleep(self.page_load_wait)

            b64 = self.chrome.take_screenshot(tab_info)
            if not b64:
                err = "screenshot failed"
                self.stats["errors"] += 1
                return

            screenshot = Image.open(BytesIO(base64.b64decode(b64)))
            ss_path = os.path.join(self.screenshot_dir, f"{ticker}.png")
            screenshot.save(ss_path)

            # ── Blind-scanner check: blank/unreadable chart? ─────────────────
            # Feeds the server-side quarantine counter — a laptop whose Chrome
            # got logged out must not silently "scan" tickers with zero results.
            if self._looks_blank(screenshot):
                ocr_ok = False
                self._local_bad_streak += 1
                self.logger.warning(
                    f"[{ticker}] 👁 Chart looks BLANK (bad streak: "
                    f"{self._local_bad_streak}) — check TradingView session!")
                # Upload evidence once, early in the streak
                if self._local_bad_streak == 3:
                    self._upload_screenshot_async(
                        ticker, "blind_scan", ss_path,
                        note="3 consecutive blank charts on this worker")
            else:
                self._local_bad_streak = 0

            self.fleet.progress(ticker, "ocr")
            boxes = self.detector.find_signal_boxes(screenshot)
            self.stats["links_processed"] += 1
            if boxes:
                boxes_found = True
                self.logger.info(f"[{ticker}] Found {len(boxes)} signal box candidate(s)")
                if self._pipeline_enabled:
                    self._ocr_queue.put((ticker, link, ss_path, screenshot, boxes))
                else:
                    self._process_link_ocr(ticker, link, ss_path, screenshot, boxes)

        except Exception as e:
            err = str(e)[:200]
            self.logger.error(f"[{ticker}] Chrome error: {e}")
            self.stats["errors"] += 1
        finally:
            if tab_id:
                try: self.chrome.close_tab(tab_id)
                except Exception: pass
            # Release this ticker's lease + record stats on the coordinator
            rep = self.fleet.report_scan(
                ticker_id    = tid,
                lease_token  = token,
                duration_sec = time.time() - t0,
                signal_found = boxes_found,
                error        = err,
                ocr_ok       = ocr_ok,
            )
            if rep and rep.get("quarantined"):
                self.logger.error(
                    "🟡 This worker was QUARANTINED by the coordinator — too many "
                    "bad scans. Fix Chrome/TradingView, then Resume from the dashboard.")
            self._current_ticker = ""
            time.sleep(self.tab_delay)

    def _drain_ocr_queue(self):
        if getattr(self, "_pipeline_enabled", False):
            try:
                self._ocr_queue.put(("__CYCLE_DONE__", None, None, None, None))
                self._ocr_queue.join()
            except Exception:
                pass


def _flt(v):
    if v is None:
        return None
    try:
        return float(str(v).replace(",", ""))
    except (ValueError, TypeError):
        return None


def main():
    print("=" * 60)
    print("  ORB FLEET WORKER  |  distributed scanning")
    print("=" * 60)
    try:
        worker = FleetORBWorker("orb_config.json")
    except Exception as e:
        print(f"❌ Startup error: {e}")
        import traceback; traceback.print_exc()
        return 1
    worker.start()
    return 0


if __name__ == "__main__":
    sys.exit(main())