#!/usr/bin/env python3

# This script produces a summary of channel forwarding stats
#
# - Displays PeerID, SCID, and Alias for each channel
# - Uses `clboss-recent-earnings` to limit the history considered
#
# The channels at the top of the list are good, the ones at the bottom are bad.

import argparse
from datetime import datetime, timezone
import json
import os
import subprocess
import time

from tabulate import tabulate
from wcwidth import wcswidth

from clboss.alias_cache import lookup_alias
from feemon_data import load_merged_records_by_node

def now():
    now_ts = time.time()
    return int(now_ts)


def run_lightning_cli_command(lightning_dir, network_option, command, *args):
    try:
        command = ["lightning-cli", network_option, command, *args]
        if lightning_dir:
            command = command[:2] + [f"--lightning-dir={lightning_dir}"] + command[2:]
        result = subprocess.run(command, capture_output=True, text=True, check=True)
        return json.loads(result.stdout)
    except subprocess.CalledProcessError as e:
        print(f"Command '{command}' failed with error: {e}")
    except json.JSONDecodeError as e:
        print(f"Failed to parse JSON from command '{command}': {e}")
    return None


def peer_age_days(peer):
    """Lookup the age in secs in the peer struct and convert to days"""
    return max(1, peer.get("age", 0) // 86400) # age in days, but not zero


def determine_peer_age(lightning_dir, network_option, id, peer_metrics):
    """Use multiple sources to determine the peer's age"""

    # What is the oldest history we have on record for this peer?
    oldest_history_age = None
    history = run_lightning_cli_command(
        lightning_dir, network_option, "clboss-earnings-history", id
    )['history']
    if len(history) >= 1 and history[0]['bucket_time'] != 0:
        # the oldest bucket is not the "legacy" bucket
        oldest_history_age = now() - history[0]['bucket_time']
    elif len(history) >= 2:
        # skip the legacy bucket
        oldest_history_age = now() - history[1]['bucket_time']

    # Do we have an age in the peer metrics (from CLBOSS)?  This gets reset
    # when the channel is closed and reopened to the same peer
    peer_metrics_age = None
    if id in peer_metrics:
        peer_metrics_age = peer_metrics[id].get("age", 0)

    if not oldest_history_age and not peer_metrics_age:
        return 365 * 86400
    elif oldest_history_age and peer_metrics_age:
        return max(oldest_history_age, peer_metrics_age)
    elif oldest_history_age:
        return oldest_history_age
    else:
        return peer_metrics_age

def peer_net_earnings(peer):
    earnings = peer.get("in_earnings", 0) + peer.get("out_earnings", 0)
    expenditures = peer.get("in_expenditures", 0) + peer.get("out_expenditures", 0)
    return earnings - expenditures


def pad_string(s, width):
    pad = width - wcswidth(s)
    return s + " " * pad


def calculate_fral(peer, avg_to_us_msat, op_days):
    if not avg_to_us_msat or op_days < 1:
        return None

    forwarded = peer.get("in_forwarded", 0) + peer.get("out_forwarded", 0)

    return float(forwarded) / float(avg_to_us_msat) / op_days


def calculate_tral(peer, avg_to_us_msat, op_days):
    if not avg_to_us_msat or op_days < 1:
        return None

    netearnings = peer_net_earnings(peer)

    # annualized, in basis points
    return float(netearnings) / float(avg_to_us_msat) / op_days * 365 * 100 * 100


def compute_operational_days(records, window_start_ts, now_ts):
    """Estimate operational days from feemon record count in the window.

    Feemon records arrive roughly once per hour, so count / 24
    gives the number of days the channel was operational.
    """
    count = sum(1 for rec in records
                if window_start_ts <= rec["ts"] <= now_ts)
    return count / 24.0


def compute_time_weighted_avg_balance(records, window_start_ts, now_ts):
    """Compute time-weighted average of balance_our_msat over [window_start_ts, now_ts].

    Uses sample-and-hold: each balance value holds until the next record.
    Records should be pre-filtered to the window period; any record at or
    before window_start_ts is used as the boundary balance.

    Returns (avg_msat_int, effective_start_ts) or (None, None) if no data.
    """
    if not records:
        return (None, None)

    # Extract (ts, balance) pairs, skipping records without balance data
    points = []
    for rec in records:
        bal = rec["fields"].get("balance_our_msat")
        if bal is not None:
            points.append((rec["ts"], int(bal)))

    if not points:
        return (None, None)

    # Find the boundary balance: last record at or before window_start_ts
    boundary_balance = None
    first_in_window_idx = None
    for i, (ts, bal) in enumerate(points):
        if ts <= window_start_ts:
            boundary_balance = bal
        else:
            if first_in_window_idx is None:
                first_in_window_idx = i

    if boundary_balance is not None:
        # We have data covering the window start
        effective_start = window_start_ts
        # Build the segments: start with boundary balance from window_start
        segments = []
        current_bal = boundary_balance
        current_ts = window_start_ts

        # Add segments for each point within the window
        for ts, bal in points:
            if ts <= window_start_ts:
                continue
            if ts >= now_ts:
                break
            segments.append((current_bal, ts - current_ts))
            current_bal = bal
            current_ts = ts

        # Final segment to now
        segments.append((current_bal, now_ts - current_ts))

    elif first_in_window_idx is not None:
        # No data before window start; average over [first_record_ts, now]
        first_ts, first_bal = points[first_in_window_idx]
        effective_start = first_ts
        segments = []
        current_bal = first_bal
        current_ts = first_ts

        for ts, bal in points[first_in_window_idx + 1:]:
            if ts >= now_ts:
                break
            segments.append((current_bal, ts - current_ts))
            current_bal = bal
            current_ts = ts

        segments.append((current_bal, now_ts - current_ts))
    else:
        # All points are after now_ts (shouldn't happen in practice)
        return (None, None)

    total_duration = sum(dur for _, dur in segments)
    if total_duration <= 0:
        # Single point at now_ts
        return (int(points[-1][1]), effective_start)

    integral = sum(bal * dur for bal, dur in segments)
    avg = int(integral / total_duration)
    return (avg, effective_start)


def main():
    parser = argparse.ArgumentParser(
        description="Run lightning-cli with specified network"
    )

    parser.add_argument("--mainnet", action="store_true", help="Run on mainnet")
    parser.add_argument("--testnet", action="store_true", help="Run on testnet")
    parser.add_argument("--signet", action="store_true", help="Run on signet")
    parser.add_argument("--regtest", action="store_true", help="Run on regtest")
    parser.add_argument("--network", help="Set the network explicitly")
    parser.add_argument("--lightning-dir", help="lightning data location")
    parser.add_argument(
        "--days",
        type=int,
        help="Number of days of history to use with clboss-recent-earnings",
    )
    parser.add_argument(
        "--db", default=None,
        help="Path to feemon sqlite database (optional; default: API only).",
    )
    parser.add_argument(
        "--sort", choices=["fral", "tral"], default="fral",
        help="Column to sort by (default: fral).",
    )

    args = parser.parse_args()

    if args.days is not None and args.days <= 0:
        parser.error("--days must be a positive integer")

    if args.db is not None and not os.path.isfile(args.db):
        parser.error(f"--db file not found: {args.db}")

    # Reconcile network option
    if args.network:
        network_option = f"--network={args.network}"
    elif args.testnet:
        network_option = "--network=testnet"
    elif args.signet:
        network_option = "--network=signet"
    elif args.regtest:
        network_option = "--network=regtest"
    else:
        network_option = (
            "--network=bitcoin"  # lightning-cli wants "bitcoin" for mainnet
        )

    if args.lightning_dir:
        lightning_dir = args.lightning_dir
        assert os.path.isdir(
            lightning_dir
        ), f'"{lightning_dir}" is not a valid directory'
    else:
        lightning_dir = None

    # Run listpeerchannels command
    listpeerchannels_data = run_lightning_cli_command(
        lightning_dir, network_option, "listpeerchannels"
    )
    if not listpeerchannels_data:
        return

    channels_data = listpeerchannels_data.get("channels", [])

    channels = {}
    peers = {}
    for channel in channels_data:
        short_channel_id = channel.get("short_channel_id")
        peer_id = channel.get("peer_id")
        to_us_msat = channel.get("to_us_msat")
        if short_channel_id and peer_id:
            channels[short_channel_id] = {
                "peerid": peer_id,
                "opener": channel.get("opener"),
                "to_us_msat": to_us_msat,
                "total_msat": channel.get("total_msat"),
            }
            if peer_id not in peers:
                peers[peer_id] = {
                    "alias": None,
                    "in_forwarded": 0,
                    "out_forwarded": 0,
                    "in_earnings": 0,
                    "out_earnings": 0,
                    "in_expenditures": 0,
                    "out_expenditures": 0,
                    "age": 0,
                }

    for peer_id in peers.keys():
        alias = lookup_alias(
            run_lightning_cli_command, lightning_dir, network_option, peer_id
        )
        peers[peer_id]["alias"] = alias

    # Get recent earnings for each peer
    if args.days:
        recent_data = run_lightning_cli_command(
            lightning_dir, network_option, "clboss-recent-earnings", str(args.days)
        )
    else:
        recent_data = run_lightning_cli_command(
            lightning_dir, network_option, "clboss-recent-earnings"
        )

    # Run clboss-status command and capture the output
    clboss_status_data = run_lightning_cli_command(
        lightning_dir, network_option, "clboss-status"
    )
    if clboss_status_data:
        peer_metrics = clboss_status_data.get("peer_metrics", {})

        recent = recent_data.get("recent", {}) if recent_data else {}

        for id in peers.keys():
            if id in recent:
                peers[id]["in_forwarded"] = recent[id].get("in_forwarded", 0)
                peers[id]["out_forwarded"] = recent[id].get("out_forwarded", 0)
                peers[id]["in_earnings"] = recent[id].get("in_earnings", 0)
                peers[id]["out_earnings"] = recent[id].get("out_earnings", 0)
                peers[id]["in_expenditures"] = recent[id].get("in_expenditures", 0)
                peers[id]["out_expenditures"] = recent[id].get("out_expenditures", 0)

                peers[id]["age"] = determine_peer_age(
                    lightning_dir, network_option, id, peer_metrics)

    # Fetch feemon data and compute time-weighted average balances
    now_ts = now()
    peer_ids = list(peers.keys())
    if args.days:
        window_start_ts = now_ts - args.days * 86400
    else:
        window_start_ts = 0

    if args.days:
        # Fetch one extra day so compute_time_weighted_avg_balance has a
        # boundary record at or before window_start_ts.
        since_dt = datetime.fromtimestamp(window_start_ts - 86400, tz=timezone.utc)
    else:
        since_dt = None

    feemon_records_by_node = load_merged_records_by_node(
        args.db, peer_ids, api_node_ids=peer_ids,
        since_dt=since_dt, before_dt=None,
        lightning_dir=lightning_dir,
        network_option=network_option,
    )

    avg_to_us_by_peer = {}
    op_days_by_peer = {}
    for peer_id in peer_ids:
        records = feemon_records_by_node.get(peer_id, [])
        avg, _ = compute_time_weighted_avg_balance(
            records, window_start_ts, now_ts
        )
        if avg is not None:
            avg_to_us_by_peer[peer_id] = avg
        op_days_by_peer[peer_id] = compute_operational_days(
            records, window_start_ts, now_ts
        )

    def channel_sort_key(cid):
        peer_id = channels[cid]["peerid"]
        op_days = op_days_by_peer.get(peer_id, 0)
        avg = avg_to_us_by_peer.get(peer_id)
        peer = peers[peer_id]
        if args.sort == "tral":
            val = calculate_tral(peer, avg, op_days)
        else:
            val = calculate_fral(peer, avg, op_days)
        sort_val = val if val is not None else float('-inf')
        return -sort_val

    sorted_channels = sorted(channels.keys(), key=channel_sort_key)

    # Prepare table data
    table_data = []
    max_alias_length = max(
        [wcswidth(peer["alias"]) for peer in peers.values() if peer["alias"]] + [5]
    )  # 5 is the length of "Alias"
    for short_channel_id in sorted_channels:
        peer_id = channels[short_channel_id]["peerid"]
        peer = peers[peer_id]
        avg = avg_to_us_by_peer.get(peer_id)
        op_days = op_days_by_peer.get(peer_id, 0)
        fral = calculate_fral(peer, avg, op_days)
        fralstr = f"{fral:.5f}" if fral is not None else "--"
        tral = calculate_tral(peer, avg, op_days)
        tralstr = f"{tral:.3f}" if tral is not None else "--"
        fwded = peer.get("in_forwarded", 0) + peer.get("out_forwarded", 0)
        netearnings = peer_net_earnings(peer)
        alias = pad_string(peer["alias"], max_alias_length)
        opener = "L" if channels[short_channel_id]["opener"] == "local" else "R"
        agestr = str(int(peer_age_days(peer)))
        opdaysstr = f"{op_days:.1f}"
        to_us_msat = f"{channels[short_channel_id]['to_us_msat']:_}"
        avg_to_us_str = f"{avg:_}" if avg is not None else "-"
        forwarded = f"{fwded:_}"
        netearningsstr = f"{netearnings:_}"
        table_data.append(
            [
                alias,
                short_channel_id,
                opener,
                to_us_msat,
                avg_to_us_str,
                agestr,
                opdaysstr,
                forwarded,
                netearningsstr,
                fralstr,
                tralstr,
            ]
        )

    # Print the table without grid
    table_str = tabulate(
        table_data,
        headers=["Alias", "SCID", "O", "curr_to_us", "avg_to_us", "Age", "OpDays", "Forwarded", "NetEarn", "FRAL**", "TRAL*"],
        tablefmt="plain",
        stralign="left",
        numalign="right",
        colalign=("left", "left", "left", "right", "right", "right", "right", "right", "right", "right", "right"),
        disable_numparse=True,
    )

    print(table_str)
    print("** Forwarding Rate on Average Liquidity (forwarded / avg balance / op days)")
    print("* Trailing Return on Average Liquidity in basis points, annualized")

if __name__ == "__main__":
    main()
