#!/usr/bin/env python3
"""
Test script for M3DB read/write functionality.
Usage: python3 test-metrics.py <BASE_URL> [USERNAME] [PASSWORD]

Examples:
    python3 test-metrics.py https://m3db.vultrlabs.dev example example
    python3 test-metrics.py http://192.168.1.100:7201
"""

import sys
import time
import random
import requests


def main():
    if len(sys.argv) < 2:
        print("Usage: python3 test-metrics.py <BASE_URL> [USERNAME] [PASSWORD]")
        print("Example: python3 test-metrics.py https://m3db.vultrlabs.dev example example")
        print("         python3 test-metrics.py http://192.168.1.100:7201")
        sys.exit(1)

    base_url = sys.argv[1].rstrip('/')
    username = sys.argv[2] if len(sys.argv) > 2 else None
    password = sys.argv[3] if len(sys.argv) > 3 else None

    # Setup auth if provided
    auth = (username, password) if username and password else None

    print(f"=== M3DB Metrics Test ===")
    print(f"URL: {base_url}")
    if auth:
        print(f"Auth: {username}:***")
    print()

    # Check coordinator health
    print("=== Health Check ===")
    health_url = f"{base_url}/health"
    try:
        resp = requests.get(health_url, auth=auth, timeout=10)
        if resp.status_code == 200:
            print(f"✓ Coordinator healthy")
        elif resp.status_code == 401:
            print(f"✗ Authentication required. Provide username and password.")
            sys.exit(1)
        else:
            print(f"✗ Coordinator unhealthy: {resp.status_code}")
            sys.exit(1)
    except requests.exceptions.RequestException as e:
        print(f"✗ Failed to connect: {e}")
        sys.exit(1)

    # Check placement
    print()
    print("=== Placement ===")
    placement_url = f"{base_url}/api/v1/services/m3db/placement"
    try:
        resp = requests.get(placement_url, auth=auth, timeout=10)
        if resp.status_code == 200:
            placement = resp.json()
            instances = placement.get("placement", {}).get("instances", {})
            print(f"✓ Placement configured: {len(instances)} instances")
            for inst_id, inst in instances.items():
                print(f"  - {inst_id}: {inst.get('endpoint', 'unknown')}")
        else:
            print(f"✗ Placement not ready: {resp.status_code}")
            print(f"  Response: {resp.text}")
    except requests.exceptions.RequestException as e:
        print(f"✗ Failed to get placement: {e}")

    # Check namespaces
    print()
    print("=== Namespaces ===")
    namespace_url = f"{base_url}/api/v1/services/m3db/namespace"
    try:
        resp = requests.get(namespace_url, auth=auth, timeout=10)
        if resp.status_code == 200:
            ns_data = resp.json()
            namespaces = ns_data.get("namespaces", {})
            print(f"✓ Namespaces configured: {len(namespaces)}")
            for ns_name in namespaces.keys():
                print(f"  - {ns_name}")
        else:
            print(f"✗ Namespaces not ready: {resp.status_code}")
    except requests.exceptions.RequestException as e:
        print(f"✗ Failed to get namespaces: {e}")

    # Query test
    print()
    print("=== Query Test ===")
    query_url = f"{base_url}/api/v1/query"
    try:
        resp = requests.get(query_url, params={"query": "up"}, auth=auth, timeout=10)
        if resp.status_code == 200:
            result = resp.json()
            status = result.get("status")
            print(f"✓ Query returned: {status}")
            data = result.get("data", {}).get("result", [])
            print(f"  Results: {len(data)} series")
        else:
            print(f"✗ Query failed: {resp.status_code}")
    except requests.exceptions.RequestException as e:
        print(f"✗ Query failed: {e}")

    # Write test using Prometheus remote_write
    print()
    print("=== Write Test ===")
    print("Writing metrics via Prometheus remote_write format...")

    try:
        import struct
        import snappy  # pip install python-snappy
    except ImportError:
        print("✗ Missing dependencies for write test")
        print("  Install with: pip install python-snappy")
        print("  Skipping write test...")
        print()
        print("=== Test complete (read-only) ===")
        return

    write_url = f"{base_url}/api/v1/prom/remote/write"

    def encode_varint(n):
        """Encode a varint"""
        result = []
        while n > 127:
            result.append((n & 0x7F) | 0x80)
            n >>= 7
        result.append(n)
        return bytes(result)

    def encode_string(field_num, s):
        """Encode a string field in protobuf"""
        data = s.encode('utf-8')
        tag = (field_num << 3) | 2
        return bytes([tag]) + encode_varint(len(data)) + data

    def encode_double(field_num, value):
        """Encode a double field in protobuf"""
        tag = (field_num << 3) | 1
        return bytes([tag]) + struct.pack('<d', value)

    def encode_int64(field_num, value):
        """Encode an int64 field in protobuf (as varint)"""
        tag = (field_num << 3) | 0
        return bytes([tag]) + encode_varint(value)

    def encode_label(name, value):
        """Encode a single Label message"""
        return encode_string(1, name) + encode_string(2, value)

    def write_metric(name, value, labels_dict):
        """Write a metric with custom labels"""
        ts_ms = int(time.time() * 1000)

        # Build all labels as repeated Label messages
        labels_data = b''

        # __name__ label first
        labels_data += bytes([0x0a]) + encode_varint(len(encode_label("__name__", name))) + encode_label("__name__", name)

        # Then custom labels
        for k, v in labels_dict.items():
            label_msg = encode_label(k, v)
            labels_data += bytes([0x0a]) + encode_varint(len(label_msg)) + label_msg

        # Build Sample (field 2 in TimeSeries)
        sample = encode_double(1, float(value)) + encode_int64(2, ts_ms)

        # Build TimeSeries
        ts_encoded = labels_data + bytes([0x12]) + encode_varint(len(sample)) + sample

        # Build WriteRequest
        write_req = bytes([0x0a]) + encode_varint(len(ts_encoded)) + ts_encoded

        # Compress with snappy
        compressed = snappy.compress(write_req)

        headers = {
            "Content-Encoding": "snappy",
            "Content-Type": "application/x-protobuf",
            "X-Prometheus-Remote-Write-Version": "0.1.0"
        }

        resp = requests.post(write_url, data=compressed, headers=headers, auth=auth, timeout=10)
        return resp.status_code

    # Write test metrics with tenant labels
    print()
    tenants = [
        {"tenant": "test-tenant", "service": "api", "env": "test"},
    ]

    ts = int(time.time())
    for labels in tenants:
        metric_name = f"test_metric_{ts}"
        metric_value = random.randint(1, 100)

        status = write_metric(metric_name, metric_value, labels)
        print(f"✓ Wrote: {metric_name} = {metric_value}")
        print(f"  Labels: tenant={labels.get('tenant')}, service={labels.get('service')}, env={labels.get('env')}")

    # Wait and query back
    time.sleep(2)

    print()
    print("=== Read Back Test ===")
    try:
        resp = requests.get(query_url, params={"query": metric_name}, auth=auth, timeout=10)
        if resp.status_code == 200:
            result = resp.json()
            data = result.get("data", {}).get("result", [])
            if data:
                print(f"✓ Metric found!")
                for series in data:
                    metric = series.get("metric", {})
                    values = series.get("values", series.get("value", []))
                    print(f"  Labels: {metric}")
                    print(f"  Values: {values}")
            else:
                print(f"✗ Metric not found (may take a moment to index)")
        else:
            print(f"✗ Query failed: {resp.status_code}")
    except requests.exceptions.RequestException as e:
        print(f"✗ Query failed: {e}")

    print()
    print("=== Multi-Tenancy Query Examples ===")
    print()
    print("Query by tenant:")
    print(f"  curl -u user:pass '{base_url}/api/v1/query?query={{tenant=\"test-tenant\"}}'")
    print()
    print("Query by service:")
    print(f"  curl -u user:pass '{base_url}/api/v1/query?query={{service=\"api\"}}'")
    print()
    print("Query by env:")
    print(f"  curl -u user:pass '{base_url}/api/v1/query?query={{env=\"test\"}}'")
    print()

    print("=== Test complete ===")


if __name__ == "__main__":
    main()