import argparse
import contextlib
import logging
import os
import socket
import subprocess
import sys
import time
from urllib.parse import urlparse

import numpy
from gpstime import GPSTimeParseAction, gpsnow
from rich.live import Live
from rich.table import Table

from . import Client, __version__, constants
from . import format as fmt
from .block import time_as_ns
from .flight import flight

logger = logging.getLogger("arrakis")


##########


parser = argparse.ArgumentParser()
parser.add_argument("--version", "-v", action="version", version=__version__)
parser.add_argument(
    "--url",
    "-u",
    type=str,
    help="initial server url",
)
parser.add_argument(
    "--ssh-proxy",
    type=str,
    metavar="HOST",
    help=(
        "create an SSH tunnel to HOST for the connection. "
        "If --url is not specified, the ARRAKIS_SERVER variable specified "
        "on the remote side will be used."
    ),
)
subparsers = parser.add_subparsers()


def add_subparser(cmd, **kwargs):
    sp = subparsers.add_parser(
        cmd.__name__,
        help=cmd.__doc__.splitlines()[0],
        description=cmd.__doc__.strip(),
        formatter_class=argparse.RawDescriptionHelpFormatter,
        **kwargs,
    )
    sp.set_defaults(func=cmd)
    return sp


##########


def _add_find_count_args(parser):
    parser.add_argument(
        "pattern",
        type=str,
        nargs="?",
        default=constants.DEFAULT_MATCH,
        help="channel pattern",
    )
    parser.add_argument(
        "--data-type",
        "--dtype",
        metavar="DTYPE",
        type=numpy.dtype,
        # action="append",
        help="data type",
    )
    parser.add_argument(
        "--min_rate",
        metavar="INT",
        type=int,
        help="minimum sample rate",
    )
    parser.add_argument(
        "--max_rate",
        metavar="INT",
        type=int,
        help="maximum sample rate",
    )
    parser.add_argument(
        "--publisher",
        metavar="ID",
        type=str,
        # action="append",
        help="publisher ID",
    )


def _add_format_args(parser):
    fgroup = parser.add_mutually_exclusive_group()
    fgroup.add_argument(
        "--table",
        action="store_const",
        dest="format",
        const="table",
        help="format output as table",
    )
    for formatter in fmt.FORMATTERS:
        fgroup.add_argument(
            f"--{formatter.name}",
            action="store_const",
            dest="format",
            const=formatter,
            help=f"format output as {formatter.desc}",
        )
    fgroup.set_defaults(format="table")
    return fgroup


def _read_channels_txt(reader):
    channels = []
    with reader as f:
        for line in f:
            if not line:
                continue
            line = line.strip()
            if line[0] == "#":
                continue
            channels.append(line)
    return channels


def _load_channel_args(channel_arg):
    channels = []
    for arg in channel_arg:
        if arg == "-":
            channels += _read_channels_txt(sys.stdin)
        else:
            channels.append(arg)
    logger.debug("channels: %s", channels)
    return channels


##################################################


def find(args):
    """find channels matching regexp pattern"""
    format_ = args.format
    del args.format
    client = Client(url=args.url)
    del args.url
    chans = client.find(**vars(args))
    if format_ == "table":
        fmt.print_rich(fmt.format_channels_rich(list(chans)))
    else:
        for chan in chans:
            print(format_(chan))


sparser = add_subparser(find, aliases=["search", "list"])
_add_find_count_args(sparser)
_add_format_args(sparser)


##########


def count(args):
    """count channels matching pattern"""
    format_ = args.format
    del args.format
    client = Client(url=args.url)
    del args.url
    cnt = client.count(**vars(args))
    if format_ == "table":
        print(cnt)
    else:
        print(format_(cnt))


sparser = add_subparser(count)
_add_find_count_args(sparser)
_add_format_args(sparser)


##########


def describe(args):
    """describe channels"""
    format_ = args.format
    del args.format
    client = Client(url=args.url)
    del args.url
    args.channels = _load_channel_args(args.channels)
    chans = client.describe(**vars(args)).values()
    if format_ == "table":
        fmt.print_rich(fmt.format_channels_rich(list(chans)))
    else:
        for chan in chans:
            print(format_(chan))


sparser = add_subparser(describe, aliases=["show"])
sparser.add_argument(
    "channels",
    nargs="+",
    help="list of channels to describe ('-' to read from stdin)",
)
_add_format_args(sparser)


##########


def stream(args):
    """stream data for channels"""
    format_ = args.format
    del args.format
    channels = _load_channel_args(args.channels)
    start = None
    end = None
    if args.start:
        start = args.start.gps()
    if args.end:
        end = args.end.gps()
    client = Client(url=args.url)

    def stream():
        next_time = None
        for i, block in enumerate(
            client.stream(
                channels,
                start=start,
                end=end,
                kafka_url=args.kafka_url,
            )
        ):
            latency = gpsnow() - block.time
            yield i, block, latency
            if next_time and block.time != next_time:
                msg = "ERROR: SKIPPED BLOCK!"
                raise RuntimeError(msg)
            next_time = block.time + block.duration

    if format_ == "table":
        with Live(auto_refresh=False) as live:
            for i, block, latency in stream():
                btable = fmt.format_block_rich(block)
                if args.latency:
                    table = Table(box=None)
                    table.add_row(btable)
                    table.add_row(f"block {i} latency: {latency} s")
                else:
                    table = btable
                live.update(table, refresh=True)
    else:
        for i, block, latency in stream():
            print(format_(block))
            if args.latency:
                print(f"block {i} latency: {latency} s", file=sys.stderr)


sparser = add_subparser(stream)
sparser.add_argument(
    "channels",
    nargs="+",
    help="list of channels to stream ('-' to read from stdin)",
)
sparser.add_argument(
    "--kafka-url",
    type=str,
    help="kafka URL",
)
sparser.add_argument(
    "--start",
    action=GPSTimeParseAction,
    help="start time (GPS or arbitrary date/time string)",
)
sparser.add_argument(
    "--end",
    action=GPSTimeParseAction,
    help="end time (GPS or arbitrary date/time string)",
)
sparser.add_argument(
    "--latency",
    action="store_true",
    help="print buffer latency to stderr",
)
_add_format_args(sparser)


##########


def fetch(args):
    """fetch data for channels"""
    format_ = args.format
    del args.format
    args.channels = _load_channel_args(args.channels)
    args.start = args.start.gps()
    if args.end:
        args.end = args.end.gps()
    elif args.duration:
        args.end = args.start + args.duration
    del args.duration
    client = Client(url=args.url)
    del args.url
    block = client.fetch(**vars(args))
    if format_ == "table":
        fmt.print_rich(fmt.format_block_rich(block))
    else:
        print(format_(block))


sparser = add_subparser(fetch)
sparser.add_argument(
    "channels",
    nargs="+",
    help="list of channels to fetch ('-' to read from stdin)",
)
sparser.add_argument(
    "--start",
    required=True,
    action=GPSTimeParseAction,
    help="start time (GPS or arbitrary date/time string)",
)
egroup = sparser.add_mutually_exclusive_group(required=True)
egroup.add_argument(
    "--end",
    action=GPSTimeParseAction,
    help="end time (GPS or arbitrary date/time string)",
)
egroup.add_argument(
    "--duration",
    type=int,
    help="duration in seconds",
)
_add_format_args(sparser)


##########


def publish(args):
    """publish values to channels

    Arguments should be channel name + generator function pairs.  The
    generator function is used to generate data for the specified
    channel and should be a sympy expression, with the 't' value used
    to indicate time, e.g. "sin(t)".  A numeric value for the
    generator function will generate a constant stream.

    """
    import sched
    import time

    from sympy import lambdify, parse_expr
    from sympy.abc import t

    from . import Publisher, SeriesBlock, Time

    if not args.list:
        if not args.channel_args or len(args.channel_args) % 2 != 0:
            parser.error("arguments must be channel name + generator function pairs")
        chan_funcs = {}
        for name, value in zip(args.channel_args[::2], args.channel_args[1::2]):
            expr = parse_expr(value)
            chan_funcs[name] = lambdify(t, expr, "numpy")

    publisher = Publisher(args.publisher_id, args.url)
    publisher.register()

    if args.list:
        for _name, channel in publisher.channels.items():
            print(fmt.formatter(channel, format="repr"))
        return

    assert publisher.stride
    stride = publisher.stride

    def _gen_data(publisher, time_ns):
        metadata = {}
        series = {}
        for name, func in chan_funcs.items():
            try:
                channel = publisher.channels[name]
            except KeyError:
                msg = f"unknown channel for publisher: {name}"
                raise ValueError(msg) from None
            samples = channel.sample_rate * stride / Time.SECONDS
            assert int(samples) == samples, f"non-integer samples: {samples}"
            # make time array for the previous block of time
            time_array = time_ns - numpy.arange(samples)
            data = numpy.array(
                numpy.broadcast_to(func(time_array), time_array.shape),
                dtype=channel.data_type,
            )
            series[name] = data
            metadata[name] = channel
        sblock = SeriesBlock(
            time_ns,
            series,
            metadata,
        )
        # logger.info("publish: %s", repr_block(sblock))
        publisher.publish(sblock)

    def gpsnow_ns():
        return time_as_ns(gpsnow())

    def sleep_ns(ns):
        s = ns / Time.SECONDS
        time.sleep(s)

    s = sched.scheduler(gpsnow_ns, sleep_ns)
    tick = int(gpsnow()) * Time.SECONDS
    with publisher:
        while True:
            tick += stride
            s.enterabs(tick, 0, _gen_data, (publisher, tick))
            s.run()


sparser = add_subparser(publish)
sparser.add_argument(
    "--publisher-id",
    "-p",
    type=str,
    required=True,
    help="publisher ID (required)",
)
lgroup = sparser.add_mutually_exclusive_group()
lgroup.add_argument(
    "--list",
    "-l",
    action="store_true",
    help="list publisher channels and exit",
)
lgroup.add_argument(
    "channel_args",
    nargs="*",
    metavar="CHANNEL FUNC",
    default=[],
    help="channel name + generator function pairs",
)


##################################################


@contextlib.contextmanager
def ssh_proxy(ssh_host, arrakis_server=None):
    # fine the remote server location if not specified
    if arrakis_server is None:
        arrakis_server = (
            subprocess.run(  # noqa S603
                ["ssh", ssh_host, "echo", "$ARRAKIS_SERVER"],
                capture_output=True,
                check=True,
            )
            .stdout.decode()
            .strip()
        )
        if not arrakis_server:
            msg = "Could not determine remote ARRAKIS_SERVER."
            raise ValueError(msg)
    parsed = urlparse(arrakis_server, scheme="grpc")

    # find an unused port.  there's always a possibility that this
    # port could be stolen between when it is grabbed below, and when
    # ssh tries to take it.
    local_host = "localhost"
    s = socket.socket()
    s.bind((local_host, 0))
    _, local_port = s.getsockname()
    s.close()

    # create ssh tunnel
    cmd = [
        "ssh",
        "-L",
        f"{local_port}:{parsed.netloc}",
        "-N",
        ssh_host,
    ]
    logger.debug("creating ssh proxy: %s", " ".join(cmd))
    proc = subprocess.Popen(cmd)  # noqa S603

    # wait for the connection to be available
    s = socket.socket()
    while True:
        try:
            time.sleep(0.1)
            s.connect(("localhost", local_port))
            s.close()
            break
        except ConnectionRefusedError:
            continue

    try:
        yield f"grpc://{local_host}:{local_port}"
    finally:
        proc.terminate()


def main():
    logger.setLevel(os.getenv("LOG_LEVEL", "DEBUG").upper())
    handler = logging.StreamHandler()
    handler.setFormatter(logging.Formatter("%(asctime)s %(message)s"))
    logger.addHandler(handler)

    args = parser.parse_args()

    if "func" not in args:
        parser.print_help()
        return

    logger.debug(args)

    func = args.func
    del args.func

    proxy: contextlib.AbstractContextManager
    if args.ssh_proxy:
        ssh_host = args.ssh_proxy
        proxy = ssh_proxy(ssh_host, args.url)
    else:
        proxy = contextlib.nullcontext()
    del args.ssh_proxy

    try:
        with proxy as url:
            args.url = url
            logger.debug("exec: %s(%s)", func.__name__, args)
            func(args)
    except flight.FlightError as e:
        msg = f"request error: {e}"
        raise SystemExit(msg) from e
    except Exception:
        raise


if __name__ == "__main__":
    import signal

    signal.signal(signal.SIGINT, signal.SIG_DFL)
    main()
