#! /usr/bin/python3

import csv
import os
import sys
import time
import requests
import argparse

# File locations
DEMO_DIR = os.path.join(os.path.dirname(__file__))
NEXMARK_SQL = os.path.join(DEMO_DIR, "nexmark.sql")

def table_sql(with_lateness):
    if with_lateness:
        lateness = "LATENESS INTERVAL 4 SECONDS"
    else:
        lateness = ""
    return f"""
CREATE TABLE person (
    id BIGINT,
    name VARCHAR,
    emailAddress VARCHAR,
    creditCard VARCHAR,
    city VARCHAR,
    state VARCHAR,
    date_time TIMESTAMP(3) NOT NULL {lateness},
    extra  VARCHAR
);
CREATE TABLE auction (
    id  BIGINT,
    itemName  VARCHAR,
    description  VARCHAR,
    initialBid  BIGINT,
    reserve  BIGINT,
    date_time  TIMESTAMP(3) NOT NULL {lateness},
    expires  TIMESTAMP(3),
    seller  BIGINT,
    category  BIGINT,
    extra  VARCHAR
);
CREATE TABLE bid (
    auction  BIGINT,
    bidder  BIGINT,
    price  BIGINT,
    channel  VARCHAR,
    url  VARCHAR,
    date_time TIMESTAMP(3) NOT NULL {lateness},
    extra  VARCHAR
);
"""

QUERY_SQL = {
    'q0': """CREATE VIEW q0 AS SELECT auction, bidder, price, date_time, extra FROM bid;""",

    'q1': """CREATE VIEW q1 AS
SELECT
    auction,
    bidder,
    0.908 * price as price, -- convert dollar to euro
    date_time,
    extra
FROM bid;""",

    'q2': """CREATE VIEW q2 AS SELECT auction, price FROM bid WHERE MOD(auction, 123) = 0;""",

    'q3': """CREATE VIEW q3 AS SELECT
    P.name, P.city, P.state, A.id
FROM
    auction AS A INNER JOIN person AS P on A.seller = P.id
WHERE
    A.category = 10 and (P.state = 'OR' OR P.state = 'ID' OR P.state = 'CA');""",

    'q4': """CREATE VIEW q4 AS
SELECT
    Q.category,
    AVG(Q.final)
FROM (
    SELECT MAX(B.price) AS final, A.category
    FROM auction A, bid B
    WHERE A.id = B.auction AND B.date_time BETWEEN A.date_time AND A.expires
    GROUP BY A.id, A.category
) Q
GROUP BY Q.category;""",

    'q7': """CREATE VIEW q7 AS
SELECT B.auction, B.price, B.bidder, B.date_time, B.extra
from bid B
JOIN (
  SELECT MAX(B1.price) AS maxprice, TUMBLE_START(B1.date_time, INTERVAL '10' SECOND) as date_time
  FROM bid B1
  GROUP BY TUMBLE(B1.date_time, INTERVAL '10' SECOND)
) B1
ON B.price = B1.maxprice
WHERE B.date_time BETWEEN B1.date_time  - INTERVAL '10' SECOND AND B1.date_time;""",

    'q8': """CREATE VIEW q8 AS
SELECT P.id, P.name, P.starttime
FROM (
  SELECT P.id, P.name,
         TUMBLE_START(P.date_time, INTERVAL '10' SECOND) AS starttime,
         TUMBLE_END(P.date_time, INTERVAL '10' SECOND) AS endtime
  FROM person P
  GROUP BY P.id, P.name, TUMBLE(P.date_time, INTERVAL '10' SECOND)
) P
JOIN (
  SELECT A.seller,
         TUMBLE_START(A.date_time, INTERVAL '10' SECOND) AS starttime,
         TUMBLE_END(A.date_time, INTERVAL '10' SECOND) AS endtime
  FROM auction A
  GROUP BY A.seller, TUMBLE(A.date_time, INTERVAL '10' SECOND)
) A
ON P.id = A.seller AND P.starttime = A.starttime AND P.endtime = A.endtime;""",

    'q9': """CREATE VIEW q9 AS
SELECT
    id, itemName, description, initialBid, reserve, date_time, expires, seller, category, extra,
    auction, bidder, price, bid_dateTime, bid_extra
FROM (
   SELECT A.*, B.auction, B.bidder, B.price, B.date_time AS bid_dateTime, B.extra AS bid_extra,
     ROW_NUMBER() OVER (PARTITION BY A.id ORDER BY B.price DESC, B.date_time ASC) AS rownum
   FROM auction A, bid B
   WHERE A.id = B.auction AND B.date_time BETWEEN A.date_time AND A.expires
)
WHERE rownum <= 1;""",

    'q10': """CREATE VIEW q10 AS -- PARTITIONED BY (dt, hm) AS
SELECT auction, bidder, price, date_time, extra, FORMAT_DATE('yyyy-MM-dd', date_time), FORMAT_DATE('HH:mm', date_time)
FROM bid;""",

    'q15': """CREATE VIEW q15 AS
SELECT
     FORMAT_DATE('yyyy-MM-dd', date_time) as 'day',
     count(*) AS total_bids,
     count(*) filter (where price < 10000) AS rank1_bids,
     count(*) filter (where price >= 10000 and price < 1000000) AS rank2_bids,
     count(*) filter (where price >= 1000000) AS rank3_bids,
     count(distinct bidder) AS total_bidders,
     count(distinct bidder) filter (where price < 10000) AS rank1_bidders,
     count(distinct bidder) filter (where price >= 10000 and price < 1000000) AS rank2_bidders,
     count(distinct bidder) filter (where price >= 1000000) AS rank3_bidders,
     count(distinct auction) AS total_auctions,
     count(distinct auction) filter (where price < 10000) AS rank1_auctions,
     count(distinct auction) filter (where price >= 10000 and price < 1000000) AS rank2_auctions,
     count(distinct auction) filter (where price >= 1000000) AS rank3_auctions
FROM bid
GROUP BY FORMAT_DATE('yyyy-MM-dd', date_time);""",

    'q17': """CREATE VIEW q17 AS
SELECT
     auction,
     format_date('yyyy-MM-dd', date_time) as 'day',
     count(*) AS total_bids,
     count(*) filter (where price < 10000) AS rank1_bids,
     count(*) filter (where price >= 10000 and price < 1000000) AS rank2_bids,
     count(*) filter (where price >= 1000000) AS rank3_bids,
     min(price) AS min_price,
     max(price) AS max_price,
     avg(price) AS avg_price,
     sum(price) AS sum_price
FROM bid
GROUP BY auction, format_date('yyyy-MM-dd', date_time);""",

    'q18': """CREATE VIEW q18 AS
SELECT auction, bidder, price, channel, url, date_time, extra
 FROM (SELECT *, ROW_NUMBER() OVER (PARTITION BY bidder, auction ORDER BY date_time DESC) AS rank_number
       FROM bid)
 WHERE rank_number <= 1;""",

    'q19': """CREATE VIEW q19 AS
SELECT * FROM
(SELECT *, ROW_NUMBER() OVER (PARTITION BY auction ORDER BY price DESC) AS rank_number FROM bid)
WHERE rank_number <= 10;""",

    'q20': """CREATE VIEW q20 AS
SELECT
    auction, bidder, price, channel, url, B.date_time, B.extra,
    itemName, description, initialBid, reserve, A.date_time as AdateTime, expires, seller, category, A.extra as Aextra
FROM
    bid AS B INNER JOIN auction AS A on B.auction = A.id
WHERE A.category = 10;""",
}

def sort_queries(queries):
    return sorted(queries, key=lambda q: int(q[1:]))
    
def parse_queries(arg):
    if arg is not None:
        queries = set()
        for s in arg:
            for q in s.split(','):
                q = q.lower()
                if q == 'all':
                    queries = set(QUERY_SQL.keys())
                elif q in QUERY_SQL:
                    queries.add(q)
                elif q != '':
                    sys.stderr.write(f'unknown query {q}\n')
                    sys.exit(1)
    else:
        queries = set(QUERY_SQL.keys())

    if len(queries) == 0:
        sys.stderr.write('no queries specified\n')
        sys.exit(1)

    return queries

def add_connector(connector_name, relation_name, is_input):
    transport_type = "kafka_" + ("input" if is_input else "output")
    json = {
        "description": "",
        "config": {
            "transport": {
                "name": transport_type,
                "config": {
                    "auto.offset.reset": "earliest",
                    "bootstrap.servers": kafka_broker,
                    "enable.ssl.certificate.verification": "true",
                    "sasl.mechanism": "PLAIN",
                    "security.protocol": "PLAINTEXT",
                }
            },
            "format": {
                "name": "csv",
                "config": {}
            }
        }
    }
    config = json["config"]["transport"]["config"]
    if is_input:
        config["enable.partition.eof"] = "true"
        config["topics"] = [connector_name]
    else:
        config["topic"] = connector_name
    requests.put(f"{api_url}/v0/connectors/{connector_name}", json=json).raise_for_status()
    return {
        "connector_name": connector_name,
        "is_input": is_input,
        "name": connector_name,
        "relation_name": relation_name,
    }
    
def add_input_connector(connector_name, relation_name):
    return add_connector(connector_name, relation_name, True)

def add_output_connector(connector_name, relation_name):
    return add_connector(connector_name, relation_name, False)

def stop_pipeline(pipeline_name, wait):
    requests.post(f"{api_url}/v0/pipelines/{pipeline_name}/shutdown").raise_for_status()
    if wait:
        return wait_for_status(pipeline_name, "Shutdown")

def start_pipeline(pipeline_name, wait):
    requests.post(f"{api_url}/v0/pipelines/{pipeline_name}/start").raise_for_status()
    if wait:
        return wait_for_status(pipeline_name, "Running")

def wait_for_status(pipeline_name, status):
    start = time.time()
    while requests.get(f"{api_url}/v0/pipelines/{pipeline_name}").json()["state"]["current_status"] != status:
        time.sleep(.1)
    return time.time() - start

def write_results(results, outfile):
    writer = csv.writer(outfile)
    writer.writerow(['when', 'runner', 'mode', 'language', 'name', 'num_cores', 'num_events', 'elapsed'])
    writer.writerows(results)
    

def main():
    # Command-line arguments
    parser = argparse.ArgumentParser(
        description='Nexmark benchmark demo'
    )
    parser.add_argument("--api-url", required=True, help="Feldera API URL (e.g., http://localhost:8080 )")
    parser.add_argument("--kafka-broker", required=True, help="Kafka broker (e.g., localhost:9092 )")
    parser.add_argument("--cores", type=int, help="Number of cores to use for workers (default: 16)")
    parser.add_argument('--lateness', action=argparse.BooleanOptionalAction, help='whether to use lateness for GC to save memory (default: --lateness)')
    parser.add_argument('--merge', action=argparse.BooleanOptionalAction, help='whether to merge all the queries into one program (default: --no-lateness)')
    parser.add_argument('--storage', action=argparse.BooleanOptionalAction, help='whether to enable storage (default: --no-storage)')
    parser.add_argument('--min-storage-rows', type=int, help='If storage is enabled, the minimum number of rows to write a batch to storage.')
    parser.add_argument('--query', action='append', help='queries to run (by default, all queries), specify one or more of: ' + ','.join(sort_queries(QUERY_SQL.keys())))
    parser.add_argument('--input-topic-suffix', help='suffix to apply to input topic names (by default, "")')
    parser.add_argument('--csv', help='File to write results in .csv format')
    parser.set_defaults(lateness=True, merge=False, storage=False, cores=16)
    
    global api_url, kafka_broker
    api_url = parser.parse_args().api_url
    kafka_broker = parser.parse_args().kafka_broker
    with_lateness = parser.parse_args().lateness
    merge = parser.parse_args().merge
    queries = sort_queries(parse_queries(parser.parse_args().query))
    cores = int(parser.parse_args().cores)
    storage = parser.parse_args().storage
    min_storage_rows = parser.parse_args().min_storage_rows
    if min_storage_rows is not None:
        min_storage_rows = int(min_storage_rows)
    suffix = parser.parse_args().input_topic_suffix or ''
    csvfile = parser.parse_args().csv

    output_connector_names = queries
    if merge and len(queries) > 1:
        merged_name = ','.join(queries)
        QUERY_SQL[merged_name] = '\n'.join([QUERY_SQL[q] for q in queries])
        queries = [merged_name]

    when = time.strftime('%Y-%m-%d %H:%M:%S', time.gmtime(time.time()))

    for program_name in queries:
        # Create program
        program_sql = table_sql(with_lateness) + QUERY_SQL[program_name]
        response = requests.put(f"{api_url}/v0/programs/{program_name}", json={
            "description": f"Nexmark benchmark: {program_name}",
            "code": program_sql
        })
        response.raise_for_status()
        program_version = response.json()["version"]

        # Compile program
        requests.post(f"{api_url}/v0/programs/{program_name}/compile", json={"version": program_version}).raise_for_status()
        
    print(f"Compiling program(s)...")
    for program_name in queries:
        while True:
            status = requests.get(f"{api_url}/v0/programs/{program_name}").json()["status"]
            print(f"Program {program_name} status: {status}")
            if status == "Success":
                break
            elif status != "Pending" and status != "CompilingRust" and status != "CompilingSql":
                raise RuntimeError(f"Failed program compilation with status {status}")
            time.sleep(5)

    input_connectors = [add_input_connector(s + suffix, s) for s in ("auction", "bid", "person")]
    output_connectors = {}
    for name in output_connector_names:
        output_connectors[name] = add_output_connector(name, name)

    # Create pipelines
    print("Creating pipeline(s)...")
    for program_name in queries:
        pipeline_name = program_name
        requests.put(f"{api_url}/v0/pipelines/{pipeline_name}", json={
            "description": "",
            "config": {"workers": cores, "storage": storage, "min_storage_rows": min_storage_rows},
            "program_name": program_name,
            "connectors": input_connectors + [output_connectors[s] for s in program_name.split(',')],
        }).raise_for_status()

    # Stop pipelines
    print("Stopping pipeline(s)...")
    for pipeline_name in queries:
        stop_pipeline(pipeline_name, False)
    for pipeline_name in queries:
        stop_pipeline(pipeline_name, True)

    # Run the pipelines
    results = []
    for pipeline_name in queries:
        start = time.time()

        # Start pipeline
        elapsed = start_pipeline(pipeline_name, True)
        print(f"Started pipeline {pipeline_name} in {elapsed:.1f} s")

        # Wait till the pipeline is completed
        start = time.time()
        last_processed = 0
        while True:
            stats = requests.get(f"{api_url}/v0/pipelines/{pipeline_name}/stats").json()
            elapsed = time.time() - start
            processed = stats["global_metrics"]["total_processed_records"]
            if processed > last_processed:
                before, after = ('\r', '') if os.isatty(1) else ('', '\n')
                sys.stdout.write(f"{before}Pipeline {pipeline_name} processed {processed} records in {elapsed:.1f} seconds{after}")
                last_processed = processed
            if stats["global_metrics"]["pipeline_complete"]:
                break
            time.sleep(.1)
        if os.isatty(1):
            print()
        elapsed = "{:.1f}".format(time.time() - start)
        print(f"Pipeline {pipeline_name} completed in {elapsed} s")

        results += [[when, "feldera", "stream", "sql", pipeline_name, cores, last_processed, elapsed]]

        # Start pipeline
        elapsed = stop_pipeline(pipeline_name, True)
        print(f"Stopped pipeline {pipeline_name} in {elapsed:.1f} s")

    write_results(results, sys.stdout)
    if csvfile is not None:
        write_results(results, open(csvfile, 'w', newline=''))


if __name__ == "__main__":
    main()