"""Quick test harness for DetectorServer with VATEX backend.

Usage:
  python -m tools.test_detector_vatex \
    --image data/example.jpeg \
    --prompt "person on left side" \
    --config models/VATEX/configs/refcoco/swin/swin_large.yaml

Requirements:
  - VATEX_DIR env pointing to models/VATEX (or override with --vatex-dir)
  - conda env 'vatex' available (or override with VATEX_CONDA_ENV / --env)
  - weights referenced in the config available

It sets DETECTOR_BACKEND=vatex and calls the MCP detector server tool over stdio.
"""
from __future__ import annotations

import argparse
import asyncio
import json
import os
import sys
from pathlib import Path
from pprint import pformat

try:
    from mcp.client.session import ClientSession  # type: ignore
    from mcp.client.stdio import stdio_client, StdioServerParameters  # type: ignore
except Exception:
    ClientSession = None  # type: ignore
    stdio_client = None  # type: ignore
    StdioServerParameters = None  # type: ignore


def _best_effort(res):
    # Direct dict (common case)
    if isinstance(res, dict):
        return res
    # MCP SDK object style
    content = getattr(res, 'content', None)
    if not content:
        return {}
    for c in content:
        # dict style part
        if isinstance(c, dict):
            ctype = c.get('type')
            if ctype in ('json', 'application/json'):
                data = c.get('json') or c.get('data') or c.get('value')
                if isinstance(data, dict):
                    return data
            if ctype == 'text' and isinstance(c.get('text'), str):
                txt = c['text']
                try:
                    parsed = json.loads(txt)
                    if isinstance(parsed, dict):
                        return parsed
                except Exception:
                    continue
        else:
            ctype = getattr(c, 'type', None)
            if ctype in ('json', 'application/json'):
                data = getattr(c, 'json', None) or getattr(c, 'data', None) or getattr(c, 'value', None)
                if isinstance(data, dict):
                    return data
            if ctype == 'text':
                txt = getattr(c, 'text', None)
                if isinstance(txt, str):
                    try:
                        parsed = json.loads(txt)
                        if isinstance(parsed, dict):
                            return parsed
                    except Exception:
                        continue
    return {}


def parse_args():
    p = argparse.ArgumentParser(description='Test detector server with VATEX backend')
    p.add_argument('--image', required=True)
    p.add_argument('--prompt', required=True)
    p.add_argument('--config', help='VATEX config path (sets VATEX_CONFIG)')
    p.add_argument('--vatex-dir', help='Override VATEX_DIR')
    p.add_argument('--env', dest='env_name', help='Conda env name (sets VATEX_CONDA_ENV)')
    return p.parse_args()


def main():
    args = parse_args()
    if ClientSession is None:
        print(json.dumps({'error': 'MCP client SDK not installed'}))
        return 1

    img = Path(args.image)
    if not img.exists():
        print(json.dumps({'error': f'image not found: {img}'}))
        return 2

    if args.vatex_dir:
        os.environ['VATEX_DIR'] = args.vatex_dir
    if 'VATEX_DIR' not in os.environ:
        # Attempt auto-detect relative to repo root (this script lives in tools/)
        repo_root = Path(__file__).resolve().parent.parent
        auto = repo_root / 'models' / 'VATEX'
        if auto.exists():
            os.environ['VATEX_DIR'] = str(auto)
        else:
            print(json.dumps({'error': 'VATEX_DIR not set; export VATEX_DIR or pass --vatex-dir', 'attempted': str(auto)}))
            return 3
    # Resolve config path if provided (allow bare filename)
    if args.config:
        conf_candidate = Path(args.config)
        if not conf_candidate.is_absolute():
            # Try relative to current working dir first
            if not conf_candidate.exists():
                vd = Path(os.environ['VATEX_DIR'])
                # Common layout: <VATEX_DIR>/configs/... or models/VATEX/configs/...
                direct = vd / args.config
                swin_path = vd / 'configs' / 'refcoco' / 'swin' / args.config
                if direct.exists():
                    conf_candidate = direct
                elif swin_path.exists():
                    conf_candidate = swin_path
        if conf_candidate.exists():
            os.environ['VATEX_CONFIG'] = str(conf_candidate.resolve())
        else:
            print(json.dumps({'warning': f'config not found: {args.config}; using default'}))
    if args.env_name:
        os.environ['VATEX_CONDA_ENV'] = args.env_name

    os.environ['DETECTOR_BACKEND'] = 'vatex'

    async def _call():
        params = StdioServerParameters(command=sys.executable, args=['-m', 'mcp.servers.detector_server'])
        async with stdio_client(params) as (read, write):
            async with ClientSession(read, write) as session:
                await session.initialize()
                res = await session.call_tool('detect', {'det_prompt': args.prompt, 'image_path': str(img)})
                return res

    try:
        res = asyncio.run(_call())
    except Exception as e:
        print(json.dumps({'error': str(e)}))
        return 1

    payload = _best_effort(res)
    if not payload:
        # Provide raw structure for debugging
        debug_repr = getattr(res, 'content', None)
        print(json.dumps({'error': 'empty payload', 'raw': pformat(debug_repr)}, ensure_ascii=False, indent=2))
    else:
        print(json.dumps(payload, ensure_ascii=False, indent=2))
    return 0


if __name__ == '__main__':
    raise SystemExit(main())
