"""
pg_statviz - stats visualization and time series analysis
"""

__author__ = "Jimmy Angelakos"
__copyright__ = "Copyright (c) 2026 Jimmy Angelakos"
__license__ = "PostgreSQL License"

import argparse
import getpass
import logging
import numpy
from argh.decorators import arg
from dateutil.parser import isoparse
from matplotlib.pyplot import close as mpclose
from matplotlib.ticker import MaxNLocator
from pandas import DataFrame
from pg_statviz.libs import plot
from pg_statviz.libs.ai import (AI_PROVIDERS, DEFAULT_AI_PROVIDER,
                                run_chart_analysis)
from pg_statviz.libs.dbconn import dbconn
from pg_statviz.libs.html_report import finalize_module_report
from pg_statviz.libs.info import getinfo


@arg('-d', '--dbname', help="database name to analyze")
@arg('-h', '--host', metavar="HOSTNAME",
     help="database server host or socket directory")
@arg('-p', '--port', help="database server port")
@arg('-U', '--username', help="database user name")
@arg('-W', '--password', action='store_true',
     help="force password prompt (should happen automatically)")
@arg('-D', '--daterange', nargs=2, metavar=('FROM', 'TO'), type=str,
     help="date range to be analyzed in ISO 8601 format e.g. 2026-01-01T00:00"
          + " 2026-01-01T23:59")
@arg('-O', '--outputdir', help="output directory")
@arg('-A', '--ai', nargs='?', const=DEFAULT_AI_PROVIDER, default=None,
     choices=AI_PROVIDERS, metavar='PROVIDER',
     help="enable AI analysis (default provider: " + DEFAULT_AI_PROVIDER
          + "). Choices: claude (Anthropic), gemini (Google AI Studio), "
            "local (Ollama vision model).")
@arg('--info', help=argparse.SUPPRESS)
@arg('--conn', help=argparse.SUPPRESS)
def io(*, dbname=getpass.getuser(), host="/var/run/postgresql", port="5432",
       username=getpass.getuser(), password=None, daterange=[],
       outputdir=None, ai=None, info=None, conn=None):
    "run I/O analysis module"

    logging.basicConfig()
    _logger = logging.getLogger(__name__)
    _logger.setLevel(logging.INFO)

    if not conn:
        conn_details = {'dbname': dbname, 'user': username,
                        'password': getpass.getpass("Password: ") if password
                        else password, 'host': host, 'port': port}
        conn = dbconn(**conn_details)
    if not info:
        info = getinfo(conn)

    _logger.info("Running I/O analysis")

    if daterange:
        daterange = [isoparse(d) for d in daterange]
        if daterange[0] > daterange[1]:
            daterange = [daterange[1], daterange[0]]
    else:
        daterange = ['-infinity', 'now()']

    # Retrieve the snapshots from DB
    cur = conn.cursor()
    cur.execute("""SELECT io_stats, block_size, i.stats_reset, snapshot_tstamp
                   FROM pgstatviz.io i
                   JOIN pgstatviz.db USING (snapshot_tstamp)
                   WHERE snapshot_tstamp BETWEEN %s AND %s
                   ORDER BY snapshot_tstamp""",
                (daterange[0], daterange[1]))

    data = cur.fetchall()
    if not data:
        cur.execute("""SELECT
                    (current_setting('server_version_num')::int >= 160000)
                    AS version_ok""")
        versioncheck = cur.fetchone()['version_ok']
        if not versioncheck:
            _logger.warning("I/O analysis is only available from "
                            + "PostgreSQL release 16 onwards")
            return
        else:
            raise SystemExit("No pg_statviz snapshots found in this database")

    tstamps = [ts['snapshot_tstamp'] for ts in data]
    blcksz = int(data[0]['block_size'])
    iostats, iokinds = calc_iostats(data, blcksz)
    iorates = calc_iorates(data, iokinds, blcksz)

    # Build a flattened DataFrame for AI analysis
    io_df = build_io_dataframe(iostats, iokinds, tstamps)

    report_sections = []

    # Plot as many of each I/O kinds we have per snapshot
    plt, fig, splt1, splt2 = plot.setupdouble()
    plt.suptitle(f"pg_statviz · {info['hostname']}:{port}",
                 fontweight='semibold')

    # Plot Reads
    splt1.set_title("I/O Reads")
    for iokind in iokinds:
        iobytes = []
        for snapshot in iostats:
            found = False
            for entry in snapshot:
                if {'backend_type': entry['backend_type'],
                        'object': entry['object'],
                        'context': entry['context']} == iokind:
                    found = True
                    r = entry['reads']
                    if r:
                        iobytes += round(r / 1073741824, 1 if r >= 100 else 2),
                    else:
                        found = False
            if not found:
                iobytes += 0,
        if not all(b == 0 for b in iobytes):
            # Downsample if needed
            _frame = DataFrame(data=iobytes, index=tstamps, copy=False)
            if len(tstamps) > plot.MAX_POINTS:
                q = str(round(
                    (tstamps[-1] - tstamps[0]).total_seconds()
                    / plot.MAX_POINTS, 2))
                r = _frame.resample(q + "s").mean()
            else:
                r = _frame
            splt1.plot_date(r.index, r,
                            label=f"{iokind['object']}/"
                                  if {iokind['object']} == 'temp relation'
                                  else ""
                                  f"{iokind['backend_type']}/"
                                  f"{iokind['context']}",
                            aa=True, linestyle='solid')
    splt1.set_xlabel("Timestamp", fontweight='semibold')
    splt1.set_ylabel("GB read (at time of snapshot)", fontweight='semibold')
    splt1.set_ylim(bottom=0)
    splt1.legend()

    # Plot Writes
    splt2.set_title("I/O Writes")
    for iokind in iokinds:
        iobytes = []
        for snapshot in iostats:
            found = False
            for entry in snapshot:
                if {'backend_type': entry['backend_type'],
                        'object': entry['object'],
                        'context': entry['context']} == iokind:
                    found = True
                    w = entry['writes']
                    if w:
                        iobytes += round(w / 1073741824, 1 if w >= 100 else 2),
                    else:
                        found = False
            if not found:
                iobytes += 0,
        if not all(b == 0 for b in iobytes):
            # Downsample if needed
            _frame = DataFrame(data=iobytes, index=tstamps, copy=False)
            if len(tstamps) > plot.MAX_POINTS:
                q = str(round(
                    (tstamps[-1] - tstamps[0]).total_seconds()
                    / plot.MAX_POINTS, 2))
                r = _frame.resample(q + "s").mean()
            else:
                r = _frame
            splt2.plot_date(r.index, r,
                            label=f"{iokind['object']}/"
                                  f"{iokind['backend_type']}/"
                                  f"{iokind['context']}",
                            aa=True, linestyle='solid')
    splt2.set_xlabel("Timestamp", fontweight='semibold')
    splt2.set_ylabel("GB written (at time of snapshot)",
                     fontweight='semibold')
    splt2.set_ylim(bottom=0)
    splt2.legend()

    fig.gca().yaxis.set_major_locator(MaxNLocator(integer=True))
    outfile = f"""{
        outputdir.rstrip("/") + "/" if outputdir
        else ''}pg_statviz_{info['hostname']
                            .replace("/", "-")}_{port}_io.png"""
    _logger.info(f"Saving {outfile}")
    plt.gcf().autofmt_xdate()
    fig.tight_layout()
    plt.savefig(outfile)
    run_chart_analysis(
        report_sections, ai, io_df, "I/O Statistics",
        metric_description="CUMULATIVE COUNTER — rising values are "
                           "NORMAL. I/O volume by component. "
                           "'client backend' is query work. "
                           "'checkpointer' spikes at checkpoint "
                           "intervals are normal. No warning "
                           "threshold. Default to [HEALTHY].",
        outfile=outfile,
        info=info,
    )

    # Build rate DataFrame for AI analysis
    rate_df = build_iorate_dataframe(iorates, iokinds, tstamps)

    # Plot I/O Rates
    plt, fig, splt1, splt2 = plot.setupdouble()
    plt.suptitle(f"pg_statviz · {info['hostname']}:{port}",
                 fontweight='semibold')

    # Plot Read Rates
    splt1.set_title("I/O Read Rate")
    for iokind in iokinds:
        iokindname = (f"{iokind['object']}/"
                      if {iokind['object']} == 'temp relation'
                      else ""
                      f"{iokind['backend_type']}/"
                      f"{iokind['context']}")
        if not all(numpy.isnan(v) or v == 0
                   for v in iorates['reads'][iokindname]):
            # Downsample if needed
            _frame = DataFrame(
                data=[round(v / 1048576, 1 if v >= 100 else 2)
                      for v in iorates['reads'][iokindname]],
                index=tstamps, copy=False)
            if len(tstamps) > plot.MAX_POINTS:
                q = str(round(
                    (tstamps[-1] - tstamps[0]).total_seconds()
                    / plot.MAX_POINTS, 2))
                r = _frame.resample(q + "s").mean()
            else:
                r = _frame
            splt1.plot_date(r.index, r, label=iokindname, aa=True,
                            linestyle='solid')
    splt1.set_xlabel("Timestamp", fontweight='semibold')
    splt1.set_ylabel("Avg. read rate in MB/s", fontweight='semibold')
    splt1.set_ylim(bottom=0)
    splt1.legend()

    # Plot Write Rates
    splt2.set_title("I/O Write Rate")
    for iokind in iokinds:
        iokindname = (f"{iokind['object']}/"
                      if {iokind['object']} == 'temp relation'
                      else ""
                      f"{iokind['backend_type']}/"
                      f"{iokind['context']}")
        if not all(numpy.isnan(v) or v == 0
                   for v in iorates['writes'][iokindname]):
            # Downsample if needed
            _frame = DataFrame(
                data=[round(v / 1048576, 1 if v >= 100 else 2)
                      for v in iorates['writes'][iokindname]],
                index=tstamps, copy=False)
            if len(tstamps) > plot.MAX_POINTS:
                q = str(round(
                    (tstamps[-1] - tstamps[0]).total_seconds()
                    / plot.MAX_POINTS, 2))
                r = _frame.resample(q + "s").mean()
            else:
                r = _frame
            splt2.plot_date(r.index, r, label=iokindname, aa=True,
                            linestyle='solid')
    splt2.set_xlabel("Timestamp", fontweight='semibold')
    splt2.set_ylabel("Avg. write rate in MB/s", fontweight='semibold')
    splt2.legend()
    splt2.set_ylim(bottom=0)

    plt.gcf().autofmt_xdate()
    fig.tight_layout()
    outfile = f"""{
        outputdir.rstrip("/") + "/" if outputdir
        else ''}pg_statviz_{info['hostname']
                            .replace("/", "-")}_{port}_io_rate.png"""
    _logger.info(f"Saving {outfile}")
    plt.savefig(outfile)
    run_chart_analysis(
        report_sections, ai, rate_df, "I/O Rate",
        metric_description="I/O RATES in MB/s (derived from cumulative "
                           "counters). NaN values appear on "
                           "stats_reset — IGNORE them. Read spikes "
                           "indicate cache misses. Write spikes "
                           "during checkpoints are normal. No "
                           "warning threshold (storage limits vary). "
                           "Default to [HEALTHY].",
        outfile=outfile,
        info=info,
    )

    finalize_module_report(outputdir, info, port, 'io',
                           report_sections)
    mpclose('all')


# Gather I/O stats and convert to bytes
def calc_iostats(data, blcksz=8192):
    iostats = [io['io_stats'] for io in data]
    iokinds = []
    for snapshot in iostats:
        for entry in snapshot:
            # PG18+ has read_bytes/write_bytes columns,
            # older versions need conversion
            if 'read_bytes' in entry:
                r = entry['read_bytes']
                if r:
                    entry['reads'] = int(r)
            else:
                r = entry['reads']
                if r:
                    entry['reads'] = int(r) * blcksz

            if 'write_bytes' in entry:
                w = entry['write_bytes']
                if w:
                    entry['writes'] = int(w)
            else:
                w = entry['writes']
                if w:
                    entry['writes'] = int(w) * blcksz

            iokind = {'backend_type': entry['backend_type'],
                      'object': entry['object'],
                      'context': entry['context']}
            if iokind not in iokinds:
                iokinds += iokind,
    return iostats, iokinds


# Calculate I/O rates
def calc_iorates(data, iokinds, blcksz=8192):

    rates = {}

    # I/O diff generator - yields tuple list of I/O rates in bytes/s
    def iodiff(data, iokind, rw):
        yield numpy.nan
        for i, item in enumerate(data):
            if i + 1 < len(data):
                if (data[i + 1]['stats_reset'] == data[i]['stats_reset']):
                    s = ((data[i + 1]['snapshot_tstamp']
                         - data[i]['snapshot_tstamp'])
                         .total_seconds())
                    v1, v2 = 0, 0
                    # PG18+ has read_bytes/write_bytes columns
                    rw_bytes = f"{rw[:-1]}_bytes"
                    for entry in data[i]['io_stats']:
                        if {'backend_type': entry['backend_type'],
                                'object': entry['object'],
                                'context': entry['context']} == iokind:
                            if rw_bytes in entry:
                                v1 = entry[rw_bytes]
                            elif rw in entry:
                                v1 = entry[rw]
                                if v1:
                                    v1 = int(v1) * blcksz
                    for entry in data[i + 1]['io_stats']:
                        if {'backend_type': entry['backend_type'],
                                'object': entry['object'],
                                'context': entry['context']} == iokind:
                            if rw_bytes in entry:
                                v2 = entry[rw_bytes]
                            elif rw in entry:
                                v2 = entry[rw]
                                if v2:
                                    v2 = int(v2) * blcksz
                    yield (v2 - v1) / s if v1 else numpy.nan
                else:
                    yield numpy.nan

    for rw in ['reads', 'writes']:
        rates[rw] = {}
        for iokind in iokinds:
            rates[rw][f"{iokind['object']}/"
                      if {iokind['object']} == 'temp relation'
                      else ""
                      f"{iokind['backend_type']}/"
                      f"{iokind['context']}"] = list(iodiff(data, iokind, rw))
    return rates


# Build a flattened DataFrame from I/O stats for AI analysis
def build_io_dataframe(iostats, iokinds, tstamps):
    data = {}
    for iokind in iokinds:
        kindname = (f"{iokind['object']}/"
                    if {iokind['object']} == 'temp relation'
                    else ""
                    f"{iokind['backend_type']}/"
                    f"{iokind['context']}")
        reads = []
        writes = []
        for snapshot in iostats:
            found = False
            for entry in snapshot:
                if {'backend_type': entry['backend_type'],
                        'object': entry['object'],
                        'context': entry['context']} == iokind:
                    found = True
                    r = entry.get('reads', 0) or 0
                    w = entry.get('writes', 0) or 0
                    reads.append(round(r / 1073741824, 2))
                    writes.append(round(w / 1073741824, 2))
                    break
            if not found:
                reads.append(0)
                writes.append(0)
        if not all(v == 0 for v in reads):
            data[f"{kindname}_read_GB"] = reads
        if not all(v == 0 for v in writes):
            data[f"{kindname}_write_GB"] = writes
    return DataFrame(data=data, index=tstamps, copy=False)


# Build a flattened DataFrame from I/O rates for AI analysis
def build_iorate_dataframe(iorates, iokinds, tstamps):
    data = {}
    for iokind in iokinds:
        kindname = (f"{iokind['object']}/"
                    if {iokind['object']} == 'temp relation'
                    else ""
                    f"{iokind['backend_type']}/"
                    f"{iokind['context']}")
        if kindname in iorates['reads']:
            reads = [round(v / 1048576, 2) if not numpy.isnan(v) else 0
                     for v in iorates['reads'][kindname]]
            if not all(v == 0 for v in reads):
                data[f"{kindname}_read_MBps"] = reads
        if kindname in iorates['writes']:
            writes = [round(v / 1048576, 2) if not numpy.isnan(v) else 0
                      for v in iorates['writes'][kindname]]
            if not all(v == 0 for v in writes):
                data[f"{kindname}_write_MBps"] = writes
    return DataFrame(data=data, index=tstamps, copy=False)
