k6 + InfluxDB + Grafana Setup - Automated script to visualize k6 load test results with automatic dashboard import and time range detection

convert_and_write_to_db_multithreaded.py

#!/usr/bin/env python3
"""
Multi-threaded streaming k6 JSON to InfluxDB converter.
Uses concurrent writes to maximize throughput for large files.
"""
import argparse
import sys
import os
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from queue import Queue
import threading

from influxdb import InfluxDBClient
from convert_json_line_to_point import convert_json_line_to_point


def write_batch_worker(client, batch, batch_num):
    """Worker function to write a batch to InfluxDB."""
    try:
        start = time.time()
        client.write_points(batch, batch_size=len(batch))
        duration = time.time() - start
        return batch_num, len(batch), duration, None
    except Exception as e:
        return batch_num, len(batch), 0, str(e)


def main(k6json_location, host, port, username, password, database,
         batch_size=5000, workers=8, verbose=False):
    """
    Stream JSON file and write to InfluxDB using multiple concurrent workers.

    Args:
        k6json_location: Path to k6 JSON file
        batch_size: Number of points per batch (default: 5000)
        workers: Number of concurrent write threads (default: 8)
        verbose: Print progress information
    """
    # Get file size for progress reporting
    file_size = os.path.getsize(k6json_location)
    file_size_mb = file_size / (1024 * 1024)
    file_size_gb = file_size_mb / 1024

    if verbose:
        print(f"Processing file: {k6json_location}")
        print(f"File size: {file_size_gb:.2f} GB ({file_size_mb:.1f} MB)")
        print(f"Batch size: {batch_size:,} points")
        print(f"Concurrent workers: {workers}")
        print()

    # Create InfluxDB client (thread-safe)
    client = InfluxDBClient(host, port, username, password, database)

    # Stats
    line_count = 0
    points_count = 0
    bytes_processed = 0
    total_read_time = 0
    total_write_time = 0
    batch_num = 0
    active_batches = 0

    # Current batch
    current_batch = []

    start_time = time.time()
    last_update = start_time

    try:
        # Create thread pool for concurrent writes
        with ThreadPoolExecutor(max_workers=workers) as executor:
            futures = []

            with open(k6json_location, 'r') as f:
                for line in f:
                    # Time the read/parse operation
                    read_start = time.time()
                    line_count += 1
                    bytes_processed += len(line)

                    # Convert line to InfluxDB point
                    point = convert_json_line_to_point(line)
                    if point:
                        current_batch.append(point)
                        points_count += 1

                    total_read_time += (time.time() - read_start)

                    # When batch is full, submit to thread pool
                    if len(current_batch) >= batch_size:
                        batch_num += 1

                        # BACKPRESSURE: Wait if we have too many in-flight batches
                        # This prevents unbounded memory growth
                        max_inflight = workers * 4  # 4 batches per worker = 32 total
                        while len(futures) >= max_inflight:
                            # Wait for at least one future to complete
                            completed = [f for f in futures if f.done()]
                            if completed:
                                for future in completed:
                                    b_num, b_size, duration, error = future.result()
                                    if error:
                                        print(f"\nError writing batch {b_num}: {error}", file=sys.stderr)
                                    else:
                                        total_write_time += duration
                                    futures.remove(future)
                                break
                            else:
                                time.sleep(0.1)  # Brief sleep before checking again

                        # Submit write task to thread pool
                        future = executor.submit(write_batch_worker, client, current_batch, batch_num)
                        futures.append(future)

                        # Start new batch
                        current_batch = []

                        # Clean up completed futures periodically
                        if batch_num % 10 == 0:  # Every 10 batches
                            completed = [f for f in futures if f.done()]
                            for future in completed:
                                b_num, b_size, duration, error = future.result()
                                if error:
                                    print(f"\nError writing batch {b_num}: {error}", file=sys.stderr)
                                else:
                                    total_write_time += duration
                                futures.remove(future)

                        # Update progress (single line, every 0.5 seconds)
                        now = time.time()
                        if verbose and (now - last_update) > 0.5:
                            percent = (bytes_processed / file_size) * 100
                            elapsed = now - start_time
                            rate = points_count / elapsed if elapsed > 0 else 0
                            eta_seconds = ((file_size - bytes_processed) / bytes_processed * elapsed) if bytes_processed > 0 else 0
                            eta_min = int(eta_seconds / 60)

                            # Calculate bottleneck percentage
                            total_measured = total_read_time + total_write_time
                            write_pct = (total_write_time / total_measured * 100) if total_measured > 0 else 0

                            # Single line with \r to overwrite
                            # Show futures queue size to diagnose backpressure
                            print(f"\rProgress: {percent:.1f}% | {points_count:,} pts | {rate:.0f} pts/s | ETA: {eta_min}m | Queue: {len(futures)} | Write: {write_pct:.0f}% | R:{total_read_time:.0f}s W:{total_write_time:.0f}s",
                                  end='', flush=True)
                            last_update = now

            # Write remaining points in current batch
            if current_batch:
                batch_num += 1
                future = executor.submit(write_batch_worker, client, current_batch, batch_num)
                futures.append(future)

            # Wait for all writes to complete
            if verbose:
                print(f"\r{'':100}", end='', flush=True)  # Clear line
                print(f"\rFinalizing... waiting for {len(futures)} batches to write...", end='', flush=True)

            for future in as_completed(futures):
                b_num, b_size, duration, error = future.result()
                if error:
                    print(f"\nError writing batch {b_num}: {error}", file=sys.stderr)
                else:
                    total_write_time += duration

        if verbose:
            print()  # New line after progress
            total_time = time.time() - start_time
            total_measured = total_read_time + total_write_time
            read_pct = (total_read_time / total_measured * 100) if total_measured > 0 else 0
            write_pct = (total_write_time / total_measured * 100) if total_measured > 0 else 0

            print()
            print(f"✓ Import complete!")
            print(f"  Lines processed: {line_count:,}")
            print(f"  Points written: {points_count:,}")
            print(f"  Total time: {total_time:.1f}s ({total_time/60:.1f} minutes)")
            print(f"  Throughput: {points_count/total_time:.0f} points/second")
            print()
            print(f"Bottleneck Analysis:")
            print(f"  Read/parse time: {total_read_time:.1f}s ({read_pct:.1f}%)")
            print(f"  Write time: {total_write_time:.1f}s ({write_pct:.1f}%)")
            if write_pct > 70:
                print(f"  → InfluxDB writes are the bottleneck (consider more workers or CPU)")
            elif read_pct > 70:
                print(f"  → File I/O is the bottleneck (disk speed limited)")
            else:
                print(f"  → Balanced performance")

    except Exception as e:
        if verbose:
            print()  # New line after progress
        print(f"Error during import: {e}", file=sys.stderr)
        raise


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description='Multi-threaded streaming k6 JSON to InfluxDB (optimized for large files)')
    parser.add_argument(
        'json', help='Location of the k6 JSON results file')
    parser.add_argument(
        '--host', nargs='?', default='localhost', help='Host for connecting to InfluxDB instance')
    parser.add_argument(
        '--port', nargs='?', default=8086, type=int, help='Port on which InfluxDB is running on the host')
    parser.add_argument(
        '--username', nargs='?', default='', help='Username for accessing InfluxDB database')
    parser.add_argument(
        '--password', nargs='?', default='', help='Password for accessing InfluxDB database')
    parser.add_argument(
        '--db', nargs='?', default='k6', help='InfluxDB database name')
    parser.add_argument(
        '--batch-size', type=int, default=5000, help='Number of points per batch (default: 5000)')
    parser.add_argument(
        '--workers', type=int, default=8, help='Number of concurrent write threads (default: 8)')
    parser.add_argument(
        '--verbose', '-v', action='store_true', help='Print progress information')

    args = parser.parse_args()
    main(args.json, args.host, args.port, args.username, args.password, args.db,
         args.batch_size, args.workers, args.verbose)

convert_and_write_to_db_streaming.py

#!/usr/bin/env python3
"""
Streaming version of k6 JSON to InfluxDB converter.
Processes file line-by-line instead of loading entire file into memory.
Much faster and more memory-efficient for large files.
"""
import argparse
import sys
import os
import time

from influxdb import InfluxDBClient
from convert_json_line_to_point import convert_json_line_to_point


def main(k6json_location, host, port, username, password, database, batch_size=5000, verbose=False):
    """
    Stream JSON file line-by-line and write to InfluxDB in batches.

    Args:
        k6json_location: Path to k6 JSON file
        batch_size: Number of points to accumulate before writing (default: 5000)
        verbose: Print progress information
    """
    # Get file size for progress reporting
    file_size = os.path.getsize(k6json_location)
    file_size_mb = file_size / (1024 * 1024)

    if verbose:
        print(f"Processing file: {k6json_location}")
        print(f"File size: {file_size_mb:.1f} MB")
        print(f"Batch size: {batch_size}")
        print()

    client = InfluxDBClient(host, port, username, password, database)

    points = []
    line_count = 0
    points_count = 0
    bytes_processed = 0

    # Timing metrics
    total_read_time = 0
    total_write_time = 0
    start_time = time.time()

    try:
        with open(k6json_location, 'r') as f:
            for line in f:
                # Time the read/parse operation
                read_start = time.time()
                line_count += 1
                bytes_processed += len(line)

                # Convert line to InfluxDB point
                point = convert_json_line_to_point(line)
                if point:
                    points.append(point)
                    points_count += 1
                total_read_time += (time.time() - read_start)

                # Write batch when we reach batch_size
                if len(points) >= batch_size:
                    # Time the write operation
                    write_start = time.time()
                    client.write_points(points, batch_size=batch_size)
                    total_write_time += (time.time() - write_start)
                    points = []  # Clear batch

                    if verbose:
                        percent = (bytes_processed / file_size) * 100
                        elapsed = time.time() - start_time
                        rate = points_count / elapsed if elapsed > 0 else 0
                        eta_seconds = ((file_size - bytes_processed) / bytes_processed * elapsed) if bytes_processed > 0 else 0
                        eta_min = int(eta_seconds / 60)

                        # Use \r to overwrite the same line
                        print(f"\rProgress: {percent:.1f}% | {points_count:,} points | {rate:.0f} pts/sec | ETA: {eta_min}m | Read: {total_read_time:.1f}s Write: {total_write_time:.1f}s", end='', flush=True)

            # Write remaining points
            if points:
                write_start = time.time()
                client.write_points(points, batch_size=batch_size)
                total_write_time += (time.time() - write_start)
                if verbose:
                    print(f"\rProgress: 100.0% | {points_count:,} points written                                    ", flush=True)

        if verbose:
            total_time = time.time() - start_time
            print()
            print(f"✓ Import complete!")
            print(f"  Lines processed: {line_count:,}")
            print(f"  Points written: {points_count:,}")
            print(f"  Total time: {total_time:.1f}s ({total_time/60:.1f} minutes)")
            print(f"  Read/parse time: {total_read_time:.1f}s ({total_read_time/total_time*100:.1f}%)")
            print(f"  Write time: {total_write_time:.1f}s ({total_write_time/total_time*100:.1f}%)")
            print(f"  Throughput: {points_count/total_time:.0f} points/second")

    except Exception as e:
        if verbose:
            print()  # New line after progress
        print(f"Error during import: {e}", file=sys.stderr)
        raise


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description='Stream k6 JSON results to InfluxDB (memory-efficient for large files)')
    parser.add_argument(
        'json', help='Location of the k6 JSON results file')
    parser.add_argument(
        '--host', nargs='?', default='localhost', help='Host for connecting to InfluxDB instance')
    parser.add_argument(
        '--port', nargs='?', default=8086, type=int, help='Port on which InfluxDB is running on the host')
    parser.add_argument(
        '--username', nargs='?', default='', help='Username for accessing InfluxDB database')
    parser.add_argument(
        '--password', nargs='?', default='', help='Password for accessing InfluxDB database')
    parser.add_argument(
        '--db', nargs='?', default='k6', help='InfluxDB database name')
    parser.add_argument(
        '--batch-size', type=int, default=5000, help='Number of points per batch (default: 5000)')
    parser.add_argument(
        '--verbose', '-v', action='store_true', help='Print progress information')

    args = parser.parse_args()
    main(args.json, args.host, args.port, args.username, args.password, args.db,
         args.batch_size, args.verbose)

GRAFANA_SETUP_README.md

k6 + InfluxDB + Grafana Setup Script

This script automates the setup of a local Grafana dashboard for viewing k6 load test results from JSON output files.

What It Does

1. Validates prerequisites (docker, python3, git, jq)
2. Creates a docker-compose stack with:
   • InfluxDB 1.8 (time-series database)
   • Grafana (visualization)
3. Starts the containers
4. Configures Grafana automatically:
   • Adds InfluxDB as a data source
   • Imports the K6 Dashboard (ID: 14801) - modern, community-maintained
5. Imports your k6 JSON results into InfluxDB
6. Opens Grafana in your browser automatically (macOS, Linux, Windows)

Prerequisites

Required:

• Docker (with docker-compose)
• Python 3
• Git

Optional but Recommended:

• jq (for automatic dashboard import)
  • macOS: brew install jq
  • Linux: apt-get install jq or yum install jq
  • Without jq, you'll need to manually import the dashboard

Performance & Resource Allocation

The script configures Docker containers with resource limits optimized for large k6 test imports:

InfluxDB:

• CPU: 4-8 cores (reserved 4, max 8)
• Memory: 4-8 GB (reserved 4GB, max 8GB)
• Optimized for write-heavy workloads (importing large JSON files)

Grafana:

• CPU: 0.5-2 cores (reserved 0.5, max 2)
• Memory: 512MB-2GB (reserved 512MB, max 2GB)
• Visualization only, minimal resources needed

Streaming Import:

• Uses memory-efficient line-by-line processing
• Writes in 5,000-point batches
• Typically 10-50 MB memory usage (vs 500MB+ for non-streaming)
• Shows bottleneck analysis (Read vs Write time)

Expected Performance (Multi-threaded):

• Small files (<100MB): <1 minute
• Medium files (1-5GB): ~3-8 minutes
• Large files (10-20GB): ~10-25 minutes
• Throughput: 50,000-100,000 points/second

Converter Versions:

• Multi-threaded (default): 8 concurrent workers, 2-4x faster
• Streaming (fallback): Single-threaded, memory-efficient
• Original (last resort): Loads entire file into memory, slow

Usage

1. Make the script executable (if not already):

   chmod +x setup-k6-grafana.sh

2. Run the script:

   ./setup-k6-grafana.sh

3. Follow the prompts:

   • Enter the path to your k6 JSON results file (supports multiple formats):
     • Just the filename: test-results.json (searches in script directory)
     • Relative path: ./results/test-results.json
     • Absolute path: /Users/you/projects/test-results.json
     • Home directory: ~/Downloads/test-results.json
   • Enter a directory name for the setup (default: k6-grafana-stack)

4. Wait for setup to complete (typically 2-5 minutes)

5. Grafana opens automatically in your browser!

   • Opens directly to the K6 Dashboard
   • Time range is automatically set based on your test timestamps
   • Username: admin / Password: admin (if prompted)
   • If the browser doesn't open automatically, manually navigate to http://localhost:3000

What Gets Created

k6-grafana-stack/           # Your chosen directory name
├── docker-compose.yml      # Docker services configuration
├── k6-dashboard.json       # Downloaded dashboard template
└── k6-json-to-influxdb-line-protocol/  # Conversion tool
    ├── convert_and_write_to_db.py         # Original converter
    ├── convert_and_write_to_db_streaming.py  # Streaming converter (used by script)
    └── venv/                                  # Python virtual environment

Files in This Gist

1. setup-k6-grafana.sh - Main automation script
2. GRAFANA_SETUP_README.md - This documentation
3. convert_and_write_to_db_multithreaded.py - Multi-threaded JSON converter (recommended!)
   • Uses 8 concurrent worker threads for maximum throughput
   • 2-4x faster than single-threaded version
   • Single-line progress display with bottleneck analysis
   • Copy to script directory before running
4. convert_and_write_to_db_streaming.py - Single-threaded streaming converter (fallback)
   • Memory-efficient but slower than multi-threaded version
   • Good for systems with limited resources

Smart File Finding

The script intelligently searches for your JSON file:

1. Just provide a filename: test-results.json

   • Script searches in current directory
   • Then searches in script directory (up to 2 levels deep)
   • Shows error with available files if not found

2. Relative paths work: ./results/test-results.json or ../data/test-results.json

3. Absolute paths work: /Users/you/projects/test-results.json

4. Tilde expansion: ~/Downloads/test-results.json expands to your home directory

5. If multiple files match: Script will list them and ask you to be more specific

6. Input is sanitized: Only alphanumeric, dash, underscore, dot, slash, tilde, and space allowed

Example Usage

./setup-k6-grafana.sh

Enter the path to your k6 JSON results file:
test-results.json                    # ← Just the filename!

✓ Found JSON file in script directory
✓ Using file: /Users/you/projects/k6-tests/test-results.json

Viewing Your Results

The script automatically:

1. ✅ Opens Grafana to the K6 Dashboard (modern, recommended)
2. ✅ Sets the time range to match your test data (with 5-minute buffer)
3. ✅ Configures all data sources

You should immediately see your test results!

Dashboard

The script imports:

• K6 Dashboard (14801) - Modern, community-maintained dashboard ⭐

If You Need to Adjust the Time Range

If some panels show "No data":

1. Click the time picker (top right)
2. Try "Last 6 hours" or "Last 24 hours"
3. Or set Absolute time range to match your test window exactly

Managing the Stack

Stop the services:

cd k6-grafana-stack
docker-compose down

Restart the services:

cd k6-grafana-stack
docker-compose up -d

Remove everything (including data):

cd k6-grafana-stack
docker-compose down -v

View logs:

cd k6-grafana-stack
docker-compose logs -f

Importing Additional JSON Files

After initial setup, you can import more k6 JSON files:

cd k6-grafana-stack/k6-json-to-influxdb-line-protocol

# Activate the virtual environment
source venv/bin/activate

# Import another JSON file
python3 convert_and_write_to_db.py /path/to/another-test.json \
  --host=localhost \
  --port=8086 \
  --username=admin \
  --password=admin \
  --db=k6

# Deactivate when done
deactivate

Troubleshooting

Dashboard shows no data

• Check the time range in Grafana matches your test timestamps
• Verify data was imported: docker-compose logs influxdb
• Check InfluxDB has data:

  docker exec -it k6-influxdb influx -username admin -password admin -database k6 -execute "SHOW MEASUREMENTS"

Can't connect to Grafana

• Wait a few more seconds for services to fully start
• Check if containers are running: docker-compose ps
• Check logs: docker-compose logs grafana

JSON import fails

• Verify your JSON file is in k6's NDJSON format (one JSON object per line)
• Check the Python script output for specific errors
• Ensure InfluxDB is running: curl http://localhost:8086/ping

Dashboard not imported

• If jq is not installed, manually import:
  1. Go to Grafana → Dashboards → Import
  2. Enter dashboard ID: 14801
  3. Select "InfluxDB-k6" as the data source
  4. Click Import

Port already in use

• If port 3000 or 8086 is already in use, edit docker-compose.yml:

  ports:
    - "3001:3000"  # Use 3001 instead of 3000

k6 JSON Output Format

The script expects k6 JSON output generated with:

k6 run --out json=results.json script.js

This produces newline-delimited JSON (NDJSON) where each line is a separate metric data point.

Running Future Tests with Live Streaming

For future tests, you can stream directly to InfluxDB without JSON files:

1. Keep the stack running
2. Run k6 with InfluxDB output:

   k6 run --out influxdb=http://localhost:8086/k6 script.js

3. Watch the dashboard in real-time!

Dashboard Features

The k6 Load Testing Results dashboard includes:

• Request rate (requests per second)
• Response time (p95, p99, average)
• Error rates (4xx, 5xx)
• Virtual Users over time
• HTTP request duration by percentile
• Data received/sent rates
• Custom metrics (if defined in your test)

Data Retention

By default, InfluxDB stores all data indefinitely. To configure retention:

docker exec -it k6-influxdb influx -username admin -password admin
> USE k6
> CREATE RETENTION POLICY "30_days" ON "k6" DURATION 30d REPLICATION 1 DEFAULT
> exit

Resources

• k6 Documentation
• InfluxDB Documentation
• Grafana Documentation
• K6 Dashboard #14801

Support

If you encounter issues:

1. Check the troubleshooting section above
2. Review container logs: docker-compose logs
3. Verify services are healthy: docker-compose ps
4. Check the GitHub repositories for the tools used

Credits

This script uses:

• k6-json-to-influxdb-line-protocol for JSON conversion
• Official k6 Load Testing Results dashboard by Grafana Labs

IMPORT_OPTIMIZATION_GUIDE.md

k6 JSON Import Optimization Guide

Current Performance Issue

Symptoms:

• Import starts fast (~22k pts/sec) then slows to ~1k pts/sec
• Memory grows from 50MB to 500MB+
• Python process maxes out one CPU core (99%)
• InfluxDB sits mostly idle (0.13% CPU)

Root Cause:

• Python's GIL (Global Interpreter Lock) limits concurrent execution
• JSON parsing (json.loads() for 50+ million lines) is CPU-intensive
• Futures list was growing unbounded (now fixed with backpressure)

Optimized Import Command

For large files (>1GB), use these optimized settings:

cd k6-grafana-stack/k6-json-to-influxdb-line-protocol
source venv/bin/activate

# Optimized settings for 10GB+ files
python3 convert_and_write_to_db_multithreaded.py ~/path/to/file.json \
  --host=localhost \
  --port=8086 \
  --username=admin \
  --password=admin \
  --db=k6 \
  --batch-size=50000 \
  --workers=4 \
  --verbose

Key Parameters:

--batch-size=50000 (default: 5000)

• Larger batches = fewer write calls = less overhead
• Reduces CPU time per point by 10x
• Trade-off: Uses more memory per batch

--workers=4 (default: 8)

• Fewer workers = less thread contention
• Better when CPU (not I/O) is the bottleneck
• Try 2, 4, or 8 and compare

Performance Comparison

Small Batches (5,000 points):

• Pro: Low memory per batch
• Con: High overhead (10,000+ write calls for 50M points)
• Result: ~1,000-5,000 pts/sec

Large Batches (50,000 points):

• Pro: Low overhead (1,000 write calls for 50M points)
• Con: ~20-50MB per batch in memory
• Result: ~20,000-50,000 pts/sec (10x faster!)

Alternative: Kill and Restart

If current import is stuck:

# 1. Kill the Python process
ps aux | grep convert_and_write | grep -v grep | awk '{print $2}' | xargs kill

# 2. Clean database
cd k6-grafana-stack
docker-compose down -v
docker-compose up -d
sleep 20

# 3. Run with optimized settings
cd k6-json-to-influxdb-line-protocol
source venv/bin/activate
python3 convert_and_write_to_db_multithreaded.py ~/path/to/LS-004.json \
  --host=localhost \
  --port=8086 \
  --username=admin \
  --password=admin \
  --db=k6 \
  --batch-size=50000 \
  --workers=4 \
  --verbose

Expected Results with Optimization

For a 14GB file with 50+ million points:

Old settings (batch=5k, workers=8):

• Time: 60-120 minutes
• Memory: 500MB-1GB
• Rate: 1,000-5,000 pts/sec

Optimized (batch=50k, workers=4):

• Time: 15-30 minutes
• Memory: 200-500MB
• Rate: 30,000-60,000 pts/sec

Monitoring Progress

The progress line shows:

Progress: 52.0% | 27,380,000 pts | 35,432 pts/s | ETA: 12m | Queue: 28 | Write: 45% | R:120s W:98s

What to watch:

• Queue should stay < 32 (backpressure working)
• pts/s should stay consistent (not degrade)
• Write % shows bottleneck (< 50% = CPU bottleneck, > 70% = DB bottleneck)

Advanced: Using orjson for Faster Parsing

Install faster JSON parser:

pip install orjson

Then modify convert_json_line_to_point.py:

import orjson  # 3-5x faster than standard json

def convert_json_line_to_point(k6json_line):
    k6json = orjson.loads(k6json_line)  # Use orjson instead
    # ... rest of code

This can increase throughput by another 2-3x!

Bottleneck Decision Tree

Is Write % > 70%?
├─ YES → InfluxDB is bottleneck
│         - Increase InfluxDB CPU/memory
│         - Reduce workers (less contention)
│         - Increase batch size
│
└─ NO → Python/CPU is bottleneck
          - Increase batch size (reduces overhead)
          - Reduce workers (less GIL contention)
          - Consider orjson for faster parsing
          - Consider splitting file and parallel import

Summary

Quick win: Use --batch-size=50000 for 10x speedup on large files!

setup-k6-grafana.sh

#!/bin/bash

set -e

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

echo -e "${BLUE}========================================${NC}"
echo -e "${BLUE}k6 + InfluxDB + Grafana Setup Script${NC}"
echo -e "${BLUE}========================================${NC}"
echo ""

# Function to check if command exists
command_exists() {
    command -v "$1" >/dev/null 2>&1
}

# Check prerequisites
echo -e "${YELLOW}Checking prerequisites...${NC}"

if ! command_exists docker; then
    echo -e "${RED}Error: docker is not installed${NC}"
    exit 1
fi

if ! command_exists docker-compose && ! docker compose version >/dev/null 2>&1; then
    echo -e "${RED}Error: docker-compose is not installed${NC}"
    exit 1
fi

if ! command_exists python3; then
    echo -e "${RED}Error: python3 is not installed${NC}"
    exit 1
fi

if ! command_exists git; then
    echo -e "${RED}Error: git is not installed${NC}"
    exit 1
fi

if ! command_exists jq; then
    echo -e "${YELLOW}Warning: jq is not installed. Dashboard import may fail.${NC}"
    echo -e "${YELLOW}Install jq with: brew install jq (macOS) or apt-get install jq (Linux)${NC}"
    JQ_AVAILABLE=false
else
    JQ_AVAILABLE=true
fi

echo -e "${GREEN}✓ All prerequisites met${NC}"
echo ""

# Get the script directory
SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"

# Prompt for JSON file
echo -e "${YELLOW}Enter the path to your k6 JSON results file:${NC}"
echo -e "${YELLOW}(You can use a filename, relative path, absolute path, or ~/path)${NC}"
read -e -r JSON_FILE_INPUT

# Sanitize input - remove any dangerous characters
# Allow: alphanumeric, dash, underscore, dot, slash, tilde, space
JSON_FILE_INPUT=$(echo "$JSON_FILE_INPUT" | sed 's/[^a-zA-Z0-9._/~[:space:]-]//g')

# Trim leading/trailing whitespace
JSON_FILE_INPUT=$(echo "$JSON_FILE_INPUT" | sed -e 's/^[[:space:]]*//' -e 's/[[:space:]]*$//')

if [ -z "$JSON_FILE_INPUT" ]; then
    echo -e "${RED}Error: No file path provided${NC}"
    exit 1
fi

# Expand tilde to home directory
JSON_FILE_INPUT="${JSON_FILE_INPUT/#\~/$HOME}"

# Try to find the file in multiple locations
JSON_FILE=""

# 1. Try as absolute or relative path from current directory
if [ -f "$JSON_FILE_INPUT" ]; then
    JSON_FILE="$JSON_FILE_INPUT"
    echo -e "${GREEN}✓ Found JSON file in current directory${NC}"
# 2. Try relative to script directory
elif [ -f "$SCRIPT_DIR/$JSON_FILE_INPUT" ]; then
    JSON_FILE="$SCRIPT_DIR/$JSON_FILE_INPUT"
    echo -e "${GREEN}✓ Found JSON file in script directory${NC}"
# 3. Check if it's just a filename and search for it in script directory
elif [[ "$JSON_FILE_INPUT" != *"/"* ]]; then
    # It's just a filename, try to find it
    FOUND_FILES=$(find "$SCRIPT_DIR" -maxdepth 2 -type f -name "$JSON_FILE_INPUT" 2>/dev/null)
    FILE_COUNT=$(echo "$FOUND_FILES" | grep -c .)

    if [ -n "$FOUND_FILES" ] && [ "$FILE_COUNT" -eq 1 ]; then
        JSON_FILE="$FOUND_FILES"
        echo -e "${GREEN}✓ Found JSON file: $(basename "$JSON_FILE")${NC}"
    elif [ "$FILE_COUNT" -gt 1 ]; then
        echo -e "${RED}Error: Multiple files found matching '$JSON_FILE_INPUT':${NC}"
        echo "$FOUND_FILES"
        echo -e "${YELLOW}Please provide a more specific path${NC}"
        exit 1
    fi
fi

# If still not found, show error with helpful message
if [ -z "$JSON_FILE" ] || [ ! -f "$JSON_FILE" ]; then
    echo -e "${RED}Error: File not found: $JSON_FILE_INPUT${NC}"
    echo ""
    echo -e "${YELLOW}Searched in:${NC}"
    echo -e "  • Current directory: $(pwd)"
    echo -e "  • Script directory: $SCRIPT_DIR"
    echo ""
    echo -e "${YELLOW}Available JSON files in script directory:${NC}"
    find "$SCRIPT_DIR" -maxdepth 2 -type f -name "*.json" 2>/dev/null | head -10 | while read -r file; do
        echo -e "  • $(basename "$file")"
    done
    exit 1
fi

# Convert to absolute path
JSON_FILE=$(cd "$(dirname "$JSON_FILE")" && pwd)/$(basename "$JSON_FILE")
echo -e "${GREEN}✓ Using file: $JSON_FILE${NC}"
echo ""

# Ask for project directory
echo -e "${YELLOW}Enter directory name for the setup (default: k6-grafana-stack):${NC}"
read -r PROJECT_DIR
PROJECT_DIR=${PROJECT_DIR:-k6-grafana-stack}

# Check if directory exists and clean up any existing setup
if [ -d "$PROJECT_DIR" ]; then
    echo -e "${YELLOW}Found existing setup directory. Cleaning up...${NC}"
    cd "$PROJECT_DIR"

    # Stop and remove any running containers
    if [ -f "docker-compose.yml" ]; then
        docker-compose down -v > /dev/null 2>&1
        echo -e "${GREEN}✓ Stopped and removed existing containers${NC}"
    fi

    # Remove old converter directory if it exists
    if [ -d "k6-json-to-influxdb-line-protocol" ]; then
        rm -rf k6-json-to-influxdb-line-protocol
    fi

    cd ..
    echo -e "${GREEN}✓ Cleaned up existing setup${NC}"
else
    echo -e "${YELLOW}Creating new setup directory...${NC}"
fi

# Create/recreate project directory
mkdir -p "$PROJECT_DIR"
cd "$PROJECT_DIR"
echo -e "${GREEN}✓ Project directory ready: $PROJECT_DIR${NC}"
echo ""

# Create docker-compose.yml
echo -e "${YELLOW}Creating docker-compose.yml...${NC}"
cat > docker-compose.yml <<'EOF'
services:
  influxdb:
    image: influxdb:1.8
    container_name: k6-influxdb
    ports:
      - "8086:8086"
    environment:
      - INFLUXDB_DB=k6
      - INFLUXDB_ADMIN_USER=admin
      - INFLUXDB_ADMIN_PASSWORD=admin
      - INFLUXDB_HTTP_AUTH_ENABLED=true
    volumes:
      - influxdb-data:/var/lib/influxdb
    networks:
      - k6-network
    deploy:
      resources:
        limits:
          cpus: '8.0'
          memory: 8G
        reservations:
          cpus: '4.0'
          memory: 4G

  grafana:
    image: grafana/grafana:latest
    container_name: k6-grafana
    ports:
      - "3000:3000"
    environment:
      - GF_SECURITY_ADMIN_USER=admin
      - GF_SECURITY_ADMIN_PASSWORD=admin
      - GF_AUTH_ANONYMOUS_ENABLED=true
      - GF_AUTH_ANONYMOUS_ORG_ROLE=Viewer
    volumes:
      - grafana-data:/var/lib/grafana
    depends_on:
      - influxdb
    networks:
      - k6-network
    deploy:
      resources:
        limits:
          cpus: '2.0'
          memory: 2G
        reservations:
          cpus: '0.5'
          memory: 512M

volumes:
  influxdb-data:
  grafana-data:

networks:
  k6-network:
    driver: bridge
EOF

echo -e "${GREEN}✓ docker-compose.yml created${NC}"
echo ""

# Start Docker containers
echo -e "${YELLOW}Starting Docker containers...${NC}"
docker-compose up -d

echo -e "${GREEN}✓ Containers started${NC}"
echo ""

# Wait for services to be ready
echo -e "${YELLOW}Waiting for services to be ready...${NC}"
sleep 10

# Wait for InfluxDB
echo -n "Waiting for InfluxDB"
INFLUX_READY=false
for i in {1..30}; do
    if curl -s --max-time 2 http://localhost:8086/ping > /dev/null 2>&1; then
        echo -e " ${GREEN}✓${NC}"
        INFLUX_READY=true
        break
    fi
    echo -n "."
    sleep 2
done

if [ "$INFLUX_READY" = false ]; then
    echo -e " ${RED}✗${NC}"
    echo -e "${RED}Error: InfluxDB did not start properly${NC}"
    echo -e "${YELLOW}Check logs with: cd $PROJECT_DIR && docker-compose logs influxdb${NC}"
    exit 1
fi

# Wait for Grafana
echo -n "Waiting for Grafana"
GRAFANA_READY=false
for i in {1..60}; do
    if curl -s --max-time 2 http://localhost:3000/api/health > /dev/null 2>&1; then
        echo -e " ${GREEN}✓${NC}"
        GRAFANA_READY=true
        break
    fi
    echo -n "."
    sleep 2

    # Check if we've been waiting too long
    if [ $i -eq 30 ]; then
        echo ""
        echo -e "${YELLOW}⚠ Grafana is taking longer than expected to start...${NC}"
        echo -n "Still waiting"
    fi
done

if [ "$GRAFANA_READY" = false ]; then
    echo -e " ${RED}✗${NC}"
    echo -e "${RED}Error: Grafana did not start properly${NC}"
    echo -e "${YELLOW}Check logs with: cd $PROJECT_DIR && docker-compose logs grafana${NC}"
    exit 1
fi

echo ""

# Configure Grafana data source via API
echo -e "${YELLOW}Configuring InfluxDB data source in Grafana...${NC}"

# Wait a bit more for Grafana to fully initialize
sleep 5

# Try to configure data source with timeout and error capture
DS_RESPONSE=$(curl -s --max-time 10 -w "\n%{http_code}" -X POST \
  -H "Content-Type: application/json" \
  -d '{
    "name": "InfluxDB-k6",
    "type": "influxdb",
    "url": "http://influxdb:8086",
    "access": "proxy",
    "database": "k6",
    "basicAuth": false,
    "isDefault": true,
    "jsonData": {
      "httpMode": "GET"
    },
    "secureJsonData": {
      "password": "admin"
    },
    "user": "admin"
  }' \
  http://admin:admin@localhost:3000/api/datasources 2>&1)

HTTP_CODE=$(echo "$DS_RESPONSE" | tail -n1)
DS_BODY=$(echo "$DS_RESPONSE" | head -n-1)

if [ "$HTTP_CODE" = "200" ] || [ "$HTTP_CODE" = "201" ]; then
    echo -e "${GREEN}✓ Data source configured${NC}"
elif [ "$HTTP_CODE" = "409" ]; then
    echo -e "${YELLOW}⚠ Data source already exists (this is fine)${NC}"
else
    echo -e "${YELLOW}⚠ Data source configuration returned HTTP $HTTP_CODE${NC}"
    echo -e "${YELLOW}  You may need to manually configure it in Grafana${NC}"
    if [ -n "$DS_BODY" ]; then
        echo -e "${YELLOW}  Response: $DS_BODY${NC}"
    fi
fi
echo ""

# Import k6 Grafana dashboard (14801 - modern version)
echo -e "${YELLOW}Importing K6 Dashboard (ID: 14801)...${NC}"

if [ "$JQ_AVAILABLE" = true ]; then
    # Download dashboard JSON
    if curl -s --max-time 10 https://grafana.com/api/dashboards/14801/revisions/1/download -o k6-dashboard.json 2>&1; then
        # Check if file was downloaded
        if [ -s k6-dashboard.json ]; then
            # Create proper import payload using jq
            # Dashboard expects DS_DUMMY as the input name
            IMPORT_PAYLOAD=$(jq -n --argjson dashboard "$(cat k6-dashboard.json)" '{
              dashboard: $dashboard,
              overwrite: true,
              inputs: [
                {
                  name: "DS_DUMMY",
                  type: "datasource",
                  pluginId: "influxdb",
                  value: "InfluxDB-k6"
                }
              ]
            }' 2>&1)

            if [ $? -eq 0 ]; then
                # Import dashboard via API
                IMPORT_RESPONSE=$(curl -s --max-time 10 -w "\n%{http_code}" -X POST \
                  -H "Content-Type: application/json" \
                  -d "$IMPORT_PAYLOAD" \
                  http://admin:admin@localhost:3000/api/dashboards/import 2>&1)

                IMPORT_CODE=$(echo "$IMPORT_RESPONSE" | tail -n1)
                IMPORT_BODY=$(echo "$IMPORT_RESPONSE" | head -n-1)

                if [ "$IMPORT_CODE" = "200" ] || [ "$IMPORT_CODE" = "201" ]; then
                    echo -e "${GREEN}✓ K6 Dashboard imported${NC}"
                    # Extract dashboard URL if available
                    DASH_URL=$(echo "$IMPORT_BODY" | jq -r '.importedUrl // empty' 2>/dev/null)
                    if [ -n "$DASH_URL" ]; then
                        echo -e "${GREEN}  Dashboard URL: http://localhost:3000${DASH_URL}${NC}"
                    fi
                elif [ "$IMPORT_CODE" = "412" ]; then
                    echo -e "${YELLOW}⚠ Dashboard already exists (this is fine)${NC}"
                else
                    echo -e "${YELLOW}⚠ Dashboard import returned HTTP $IMPORT_CODE${NC}"
                    if echo "$IMPORT_BODY" | grep -q "name-exists"; then
                        echo -e "${YELLOW}  Dashboard may already exist - continuing${NC}"
                    else
                        echo -e "${YELLOW}  You can manually import dashboard ID 14801 from Grafana${NC}"
                    fi
                fi
            else
                echo -e "${YELLOW}⚠ Failed to create import payload${NC}"
                echo -e "${YELLOW}  You can manually import dashboard ID 14801 from Grafana${NC}"
            fi
        else
            echo -e "${YELLOW}⚠ Failed to download dashboard${NC}"
            echo -e "${YELLOW}  You can manually import dashboard ID 14801 from Grafana${NC}"
        fi
    else
        echo -e "${YELLOW}⚠ Failed to download dashboard from Grafana.com${NC}"
        echo -e "${YELLOW}  You can manually import dashboard ID 14801 from Grafana${NC}"
    fi
else
    echo -e "${YELLOW}⚠ Skipping automatic dashboard import (jq not installed)${NC}"
    echo -e "${YELLOW}  Manually import dashboard: Go to Grafana → Dashboards → Import → Use ID: 14801${NC}"
fi
echo ""

# Clone the k6-json-to-influxdb converter
echo -e "${YELLOW}Setting up JSON to InfluxDB converter...${NC}"

if [ ! -d "k6-json-to-influxdb-line-protocol" ]; then
    if git clone https://github.com/yooap/k6-json-to-influxdb-line-protocol.git > /dev/null 2>&1; then
        echo -e "${GREEN}✓ Converter cloned${NC}"
    else
        echo -e "${RED}✗ Failed to clone converter repository${NC}"
        echo -e "${YELLOW}Check your internet connection or try again later${NC}"
        exit 1
    fi
else
    echo -e "${GREEN}✓ Converter already exists${NC}"
fi

cd k6-json-to-influxdb-line-protocol

# Copy optimized converters if available
if [ -f "$SCRIPT_DIR/convert_and_write_to_db_multithreaded.py" ]; then
    cp "$SCRIPT_DIR/convert_and_write_to_db_multithreaded.py" .
    chmod +x convert_and_write_to_db_multithreaded.py
    echo -e "${GREEN}✓ Multi-threaded converter installed${NC}"
elif [ -f "$SCRIPT_DIR/convert_and_write_to_db_streaming.py" ]; then
    cp "$SCRIPT_DIR/convert_and_write_to_db_streaming.py" .
    chmod +x convert_and_write_to_db_streaming.py
    echo -e "${GREEN}✓ Streaming converter installed${NC}"
fi

# Create and activate virtual environment
echo -e "${YELLOW}Creating Python virtual environment...${NC}"
if python3 -m venv venv 2>&1; then
    echo -e "${GREEN}✓ Virtual environment created${NC}"
else
    echo -e "${RED}✗ Failed to create virtual environment${NC}"
    cd ..
    exit 1
fi

# Activate virtual environment
source venv/bin/activate

# Install Python dependencies
echo -e "${YELLOW}Installing Python dependencies in venv...${NC}"
if [ -f "requirements.txt" ]; then
    if python3 -m pip install -q -r requirements.txt 2>&1; then
        echo -e "${GREEN}✓ Dependencies installed${NC}"
    else
        echo -e "${YELLOW}⚠ Some dependencies may have failed to install${NC}"
    fi
else
    # Install manually if requirements.txt doesn't exist
    if python3 -m pip install -q influxdb 2>&1; then
        echo -e "${GREEN}✓ Dependencies installed (influxdb)${NC}"
    else
        echo -e "${RED}✗ Failed to install influxdb package${NC}"
        deactivate
        cd ..
        exit 1
    fi
fi
echo ""

# Import JSON data to InfluxDB using optimized converter
echo -e "${YELLOW}Importing k6 JSON data to InfluxDB...${NC}"

# Get file size for progress indication
FILE_SIZE=$(ls -lh "$JSON_FILE" | awk '{print $5}')
FILE_SIZE_BYTES=$(stat -f%z "$JSON_FILE" 2>/dev/null || stat -c%s "$JSON_FILE" 2>/dev/null)
FILE_SIZE_MB=$(echo "scale=1; $FILE_SIZE_BYTES / 1024 / 1024" | bc 2>/dev/null || echo "?")

echo -e "${BLUE}File size: ${FILE_SIZE} (${FILE_SIZE_MB} MB)${NC}"

# Select best converter
if [ -f "convert_and_write_to_db_multithreaded.py" ]; then
    CONVERTER_SCRIPT="convert_and_write_to_db_multithreaded.py"
    # Use 100k batch size - sweet spot between performance and HTTP limits
    CONVERTER_ARGS="--verbose --workers=4 --batch-size=100000"
    echo -e "${BLUE}Using multi-threaded import (4 workers, 100k batch size)${NC}"
elif [ -f "convert_and_write_to_db_streaming.py" ]; then
    CONVERTER_SCRIPT="convert_and_write_to_db_streaming.py"
    CONVERTER_ARGS="--verbose"
    echo -e "${BLUE}Using streaming import${NC}"
else
    CONVERTER_SCRIPT="convert_and_write_to_db.py"
    CONVERTER_ARGS=""
    echo -e "${BLUE}Using standard import${NC}"
fi
echo ""

# Run the import (in venv) with live progress output
if python3 $CONVERTER_SCRIPT "$JSON_FILE" \
  --host=localhost \
  --port=8086 \
  --username=admin \
  --password=admin \
  --db=k6 \
  $CONVERTER_ARGS 2>&1; then
    echo ""
    echo -e "${GREEN}✓ JSON data imported successfully${NC}"

    # Get final count
    FINAL_COUNT=$(docker exec k6-influxdb influx -username admin -password admin -database k6 -execute "SELECT COUNT(*) FROM http_req_duration" 2>/dev/null | tail -1 | awk '{print $2}')
    if [ -n "$FINAL_COUNT" ] && [ "$FINAL_COUNT" != "count_value" ]; then
        echo -e "${GREEN}  Total requests in database: ${FINAL_COUNT}${NC}"
    fi
else
    IMPORT_EXIT_CODE=$?
    echo ""
    echo -e "${RED}✗ Failed to import JSON data (exit code: $IMPORT_EXIT_CODE)${NC}"
    echo -e "${YELLOW}Possible issues:${NC}"
    echo -e "${YELLOW}  • JSON file may not be in k6 NDJSON format${NC}"
    echo -e "${YELLOW}  • InfluxDB may not be ready${NC}"
    echo -e "${YELLOW}  • Check the error messages above${NC}"
    echo ""
    echo -e "${YELLOW}Try checking InfluxDB logs:${NC}"
    echo -e "${YELLOW}  cd .. && docker-compose logs influxdb${NC}"
    deactivate
    cd ..
    exit 1
fi
echo ""

# Deactivate venv
deactivate

cd ..

# Final instructions
echo -e "${GREEN}========================================${NC}"
echo -e "${GREEN}Setup Complete! 🎉${NC}"
echo -e "${GREEN}========================================${NC}"
echo ""
echo -e "${BLUE}Services:${NC}"
echo -e "  • Grafana:   ${GREEN}http://localhost:3000${NC} (admin/admin)"
echo -e "  • InfluxDB:  ${GREEN}http://localhost:8086${NC} (admin/admin)"
echo ""
echo -e "${BLUE}Next Steps:${NC}"
echo -e "  1. Grafana opened automatically with the K6 Dashboard"
echo -e "  2. Time range is preset to match your test data"
echo -e "  3. View your test results!"
echo ""
echo -e "${YELLOW}Important:${NC}"
echo -e "  • Time range is automatically calculated from your test data"
echo -e "  • If some panels show 'No data', try adjusting the time range slightly"
echo -e "  • Dashboard data source is configured as 'InfluxDB-k6'"
echo -e "  • Dashboard: K6 Dashboard (ID: 14801) - Modern, community-maintained"
echo ""
echo -e "${BLUE}To stop the services:${NC}"
echo -e "  cd $(pwd) && docker-compose down"
echo ""
echo -e "${BLUE}To restart the services:${NC}"
echo -e "  cd $(pwd) && docker-compose up -d"
echo ""
echo -e "${BLUE}To remove everything (including data):${NC}"
echo -e "  cd $(pwd) && docker-compose down -v"
echo ""

# Extract time range from JSON file
echo -e "${YELLOW}Calculating time range from test data...${NC}"

# Get first and last timestamps from JSON
FIRST_TIME=$(grep -o '"time":"[^"]*"' "$JSON_FILE" | head -1 | sed 's/"time":"//;s/"//')
LAST_TIME=$(grep -o '"time":"[^"]*"' "$JSON_FILE" | tail -1 | sed 's/"time":"//;s/"//')

if [ -n "$FIRST_TIME" ] && [ -n "$LAST_TIME" ]; then
    # Convert to milliseconds for Grafana URL (with 5 minute buffer before/after)
    FROM_MS=$(python3 -c "
import datetime
dt = datetime.datetime.fromisoformat('$FIRST_TIME'.replace('Z', '+00:00'))
# Subtract 5 minutes for buffer
dt = dt - datetime.timedelta(minutes=5)
print(int(dt.timestamp() * 1000))
" 2>/dev/null)

    TO_MS=$(python3 -c "
import datetime
dt = datetime.datetime.fromisoformat('$LAST_TIME'.replace('Z', '+00:00'))
# Add 5 minutes for buffer
dt = dt + datetime.timedelta(minutes=5)
print(int(dt.timestamp() * 1000))
" 2>/dev/null)

    if [ -n "$FROM_MS" ] && [ -n "$TO_MS" ]; then
        echo -e "${GREEN}✓ Time range detected: $FIRST_TIME to $LAST_TIME${NC}"
        DASHBOARD_URL="http://localhost:3000/d/9lcthCWnk/k6-dashboard?from=${FROM_MS}&to=${TO_MS}"
    else
        echo -e "${YELLOW}⚠ Could not calculate time range, using default${NC}"
        DASHBOARD_URL="http://localhost:3000"
    fi
else
    echo -e "${YELLOW}⚠ Could not extract timestamps from JSON file${NC}"
    DASHBOARD_URL="http://localhost:3000"
fi

# Open Grafana in browser
echo -e "${YELLOW}Opening Grafana dashboard...${NC}"
sleep 2

# Detect OS and open browser
if [[ "$OSTYPE" == "darwin"* ]]; then
    # macOS
    open "$DASHBOARD_URL"
elif [[ "$OSTYPE" == "linux-gnu"* ]]; then
    # Linux
    if command_exists xdg-open; then
        xdg-open "$DASHBOARD_URL" 2>/dev/null
    elif command_exists gnome-open; then
        gnome-open "$DASHBOARD_URL" 2>/dev/null
    else
        echo -e "${YELLOW}⚠ Could not auto-open browser. Please manually navigate to ${DASHBOARD_URL}${NC}"
    fi
elif [[ "$OSTYPE" == "msys" || "$OSTYPE" == "cygwin" ]]; then
    # Windows (Git Bash or Cygwin)
    start "$DASHBOARD_URL"
else
    echo -e "${YELLOW}⚠ Could not auto-open browser. Please manually navigate to ${DASHBOARD_URL}${NC}"
fi

echo -e "${GREEN}✓ Browser opened with time range preset${NC}"
echo ""