#!/usr/bin/env python3
"""
Research script harness - runs implementation.py with API initialization,
stdout/stderr capture, and matplotlib figure capture.

Can be called in-process (preferred) via run() or as a subprocess for backward
compatibility.

Usage (subprocess):
    python -m dexorder.tools.research_harness <implementation_path>

Output (JSON to stdout):
    {
        "stdout": "captured user stdout",
        "stderr": "captured user stderr",
        "images": [{"format": "png", "data": "<base64>"}],
        "error": false
    }
"""

import io
import os
import base64
import json
import sys
import traceback
from pathlib import Path

# Non-interactive matplotlib backend (must be set before importing pyplot).
# Idempotent — safe to call multiple times.
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

# Ensure dexorder package is importable when run as a subprocess
sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent))


def run(impl_path: Path, item_dir: Path) -> dict:
    """
    Run a research script in-process and return captured results.

    Creates a fresh DataAPIImpl per call (thread-safe: API stored in thread-local
    via set_api() so the global API is not overwritten).

    Returns:
        dict with stdout, stderr, images, error fields
    """
    impl_path = Path(impl_path)

    if not impl_path.exists():
        return {
            "stdout": "",
            "stderr": f"Implementation file not found: {impl_path}",
            "images": [],
            "error": True,
        }

    # ---------------------------------------------------------------------------
    # Initialize a fresh API instance for this execution (thread-local)
    # ---------------------------------------------------------------------------
    try:
        import yaml

        config_path = os.environ.get("CONFIG_PATH", "/app/config/config.yaml")
        secrets_path = os.environ.get("SECRETS_PATH", "/app/config/secrets.yaml")

        config_data = {}
        secrets_data = {}
        if Path(config_path).exists():
            with open(config_path) as f:
                config_data = yaml.safe_load(f) or {}
        if Path(secrets_path).exists():
            with open(secrets_path) as f:
                secrets_data = yaml.safe_load(f) or {}

        data_cfg = config_data.get("data", {})
        iceberg_cfg = data_cfg.get("iceberg", {})
        relay_cfg = data_cfg.get("relay", {})

        from dexorder.api import set_api, API
        from dexorder.impl.charting_api_impl import ChartingAPIImpl
        from dexorder.impl.data_api_impl import DataAPIImpl

        _data_api = DataAPIImpl(
            iceberg_catalog_uri=iceberg_cfg.get("catalog_uri", "http://iceberg-catalog:8181"),
            relay_endpoint=relay_cfg.get("endpoint", "tcp://relay:5559"),
            notification_endpoint=relay_cfg.get("notification_endpoint", "tcp://relay:5558"),
            namespace=iceberg_cfg.get("namespace", "trading"),
            s3_endpoint=iceberg_cfg.get("s3_endpoint") or secrets_data.get("s3_endpoint"),
            s3_access_key=iceberg_cfg.get("s3_access_key") or secrets_data.get("s3_access_key"),
            s3_secret_key=iceberg_cfg.get("s3_secret_key") or secrets_data.get("s3_secret_key"),
            s3_region=iceberg_cfg.get("s3_region") or secrets_data.get("s3_region"),
            request_timeout=240.0,
        )
        # NOTE: We intentionally do NOT call asyncio.run(_data_api.start()) here.
        # DataAPIImpl.historical_ohlc() auto-starts on first use, which ensures the
        # ZMQ context and notification listener are created inside the user's own
        # asyncio.run() event loop — avoiding cross-loop lifecycle issues.
        # In a harness thread, set_api() stores to thread-local (not the global).
        set_api(API(charting=ChartingAPIImpl(), data=_data_api))
    except Exception as e:
        # Non-fatal — script may not use the API
        sys.stderr.write(f"WARNING: API initialization failed: {e}\n")

    # ---------------------------------------------------------------------------
    # Register custom indicators
    # ---------------------------------------------------------------------------
    try:
        from dexorder.tools.python_tools import setup_custom_indicators
        _data_dir = Path(os.environ.get("DATA_DIR", "/app/data"))
        setup_custom_indicators(_data_dir)
    except Exception as e:
        sys.stderr.write(f"WARNING: Custom indicator registration failed: {e}\n")

    # ---------------------------------------------------------------------------
    # Execute user script with captured stdout/stderr
    # ---------------------------------------------------------------------------
    stdout_buf = io.StringIO()
    stderr_buf = io.StringIO()

    # Eagerly capture figures when user scripts call plt.close() so images are
    # not lost even if the script closes figures immediately after savefig().
    captured_images: list[dict] = []

    def _capture_fig(fig) -> None:
        buf = io.BytesIO()
        fig.savefig(buf, format='png', dpi=100, bbox_inches='tight')
        buf.seek(0)
        captured_images.append({"format": "png", "data": base64.b64encode(buf.read()).decode('utf-8')})
        buf.close()

    _orig_plt_close = plt.close

    def _patched_close(fig=None):
        if fig is None:
            for fn in plt.get_fignums():
                _capture_fig(plt.figure(fn))
        elif fig == 'all':
            for fn in plt.get_fignums():
                _capture_fig(plt.figure(fn))
        else:
            try:
                _capture_fig(fig if hasattr(fig, 'savefig') else plt.figure(fig))
            except Exception:
                pass
        _orig_plt_close(fig)

    error_occurred = False
    old_stdout, old_stderr = sys.stdout, sys.stderr
    old_cwd = os.getcwd()
    sys.stdout = stdout_buf
    sys.stderr = stderr_buf
    plt.close = _patched_close

    try:
        os.chdir(impl_path.parent)
        exec(compile(impl_path.read_text(), str(impl_path), 'exec'), {})
    except Exception as e:
        print(f"ERROR: {e}", file=sys.stderr)
        traceback.print_exc(file=sys.stderr)
        error_occurred = True
    finally:
        sys.stdout = old_stdout
        sys.stderr = old_stderr
        os.chdir(old_cwd)
        plt.close = _orig_plt_close

    stdout_output = stdout_buf.getvalue()
    stderr_output = stderr_buf.getvalue()

    # ---------------------------------------------------------------------------
    # Capture any figures still open after script completion
    # ---------------------------------------------------------------------------
    images = captured_images
    if not error_occurred:
        already_seen = {img["data"] for img in images}
        for fig_num in plt.get_fignums():
            fig = plt.figure(fig_num)
            buf = io.BytesIO()
            fig.savefig(buf, format='png', dpi=100, bbox_inches='tight')
            buf.seek(0)
            data = base64.b64encode(buf.read()).decode('utf-8')
            buf.close()
            if data not in already_seen:
                images.append({"format": "png", "data": data})
    plt.close('all')

    return {
        "stdout": stdout_output,
        "stderr": stderr_output,
        "images": images,
        "error": error_occurred,
    }


def main():
    if len(sys.argv) < 2:
        print("Usage: research_harness.py <implementation_path>", file=sys.stderr)
        sys.exit(2)

    impl_path = Path(sys.argv[1])
    item_dir = impl_path.parent
    result = run(impl_path, item_dir)
    print(json.dumps(result))


if __name__ == "__main__":
    main()
