#!/usr/bin/env -S uv run --quiet --script
# /// script
# requires-python = ">=3.11"
# dependencies = [
#   "requests",
# ]
# ///
"""
Query Loki for raw OTEL JSON logs without formatting
Useful for debugging and getting exact log data
Can parse Grafana Explore URLs directly
"""

import sys
import json
import requests
from datetime import datetime, timedelta
import argparse
from urllib.parse import urlparse, parse_qs, unquote
import re

LOKI_URL = "http://localhost:3100"

def parse_grafana_time(grafana_time):
    """Parse Grafana time format like 'now-6h' into timedelta"""
    if grafana_time == 'now':
        return timedelta(0)
    
    # Match patterns like "now-6h", "now-30m", "now-2d"
    match = re.match(r'now-(\d+)([smhd])', grafana_time)
    if match:
        value = int(match.group(1))
        unit = match.group(2)
        
        if unit == 's':
            return timedelta(seconds=value)
        elif unit == 'm':
            return timedelta(minutes=value)
        elif unit == 'h':
            return timedelta(hours=value)
        elif unit == 'd':
            return timedelta(days=value)
    
    return timedelta(hours=1)

def parse_grafana_url(url):
    """Parse Grafana Explore URL and extract LogQL query and time range"""
    try:
        parsed = urlparse(url)
        query_params = parse_qs(parsed.query)
        
        if 'panes' not in query_params:
            return None, None
        
        # Decode and parse the panes JSON
        panes_json = query_params['panes'][0]
        panes_data = json.loads(panes_json)
        
        # Extract first pane (usually the only one)
        first_pane = list(panes_data.values())[0]
        
        # Extract query
        queries = first_pane.get('queries', [])
        if not queries:
            return None, None
        
        logql_query = queries[0].get('expr', '')
        
        # Extract time range
        time_range_data = first_pane.get('range', {})
        from_time = time_range_data.get('from', 'now-1h')
        
        # Convert Grafana time to our format
        time_delta = parse_grafana_time(from_time)
        
        # Convert timedelta to simple format
        total_seconds = int(time_delta.total_seconds())
        if total_seconds < 60:
            time_str = f"{total_seconds}s"
        elif total_seconds < 3600:
            time_str = f"{total_seconds // 60}m"
        elif total_seconds < 86400:
            time_str = f"{total_seconds // 3600}h"
        else:
            time_str = f"{total_seconds // 86400}d"
        
        return logql_query, time_str
    
    except Exception as e:
        print(f"Error parsing Grafana URL: {e}", file=sys.stderr)
        return None, None

def parse_time_range(time_str):
    """Parse time range string like '5m', '1h', '6h' into timedelta"""
    if not time_str:
        return timedelta(hours=1)
    
    unit = time_str[-1]
    try:
        value = int(time_str[:-1])
    except ValueError:
        return timedelta(hours=1)
    
    if unit == 's':
        return timedelta(seconds=value)
    elif unit == 'm':
        return timedelta(minutes=value)
    elif unit == 'h':
        return timedelta(hours=value)
    elif unit == 'd':
        return timedelta(days=value)
    else:
        return timedelta(hours=1)

def query_loki_raw(query, time_range='1h', limit=100):
    """Query Loki and return raw log lines"""
    end = datetime.now()
    start = end - parse_time_range(time_range)
    
    params = {
        'query': query,
        'start': int(start.timestamp() * 1e9),
        'end': int(end.timestamp() * 1e9),
        'limit': limit,
        'direction': 'backward'
    }
    
    try:
        url = f"{LOKI_URL}/loki/api/v1/query_range"
        response = requests.get(url, params=params, timeout=30)
        response.raise_for_status()
        return response.json()
    except requests.exceptions.ConnectionError:
        print("❌ Cannot connect to Loki. Is port-forward running?", file=sys.stderr)
        print("\nRun this in another terminal:", file=sys.stderr)
        print("  kubectl port-forward -n observability svc/loki-gateway 3100:80 --context=dev", file=sys.stderr)
        sys.exit(1)
    except requests.exceptions.RequestException as e:
        print(f"❌ Error querying Loki: {e}", file=sys.stderr)
        sys.exit(1)

def main():
    parser = argparse.ArgumentParser(
        description='Query Loki for raw OTEL JSON logs',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # From Grafana URL
  query-raw-logs --url 'https://telemetry.tutero.dev/grafana/explore?panes=...'
  
  # Direct LogQL query
  query-raw-logs --query '{app="resources-graphql"} |= `[latex-processor]`' --time 6h
  
  # Build query from parts
  query-raw-logs --app resources-graphql --filter "[latex-processor]" --time 2h
  query-raw-logs --app resources-graphql --severity ERROR --time 1h
        """
    )
    
    # Option 1: Grafana URL
    parser.add_argument('--url', '-u', help='Grafana Explore URL (will extract query and time range)')
    
    # Option 2: Direct LogQL query
    parser.add_argument('--query', '-q', help='Raw LogQL query string')
    
    # Option 3: Build query from parts
    parser.add_argument('--app', help='App label to filter (e.g., resources-graphql)')
    parser.add_argument('--namespace', default='learning', help='Namespace (default: learning)')
    parser.add_argument('--filter', '-f', help='Text to filter with |= operator')
    parser.add_argument('--severity', choices=['DEBUG', 'INFO', 'WARN', 'ERROR'], help='Filter by severity (requires JSON parsing)')
    
    # Common options
    parser.add_argument('--time', '-t', help='Time range (e.g., 5m, 1h, 6h) - auto-extracted from URL if not specified')
    parser.add_argument('--limit', '-l', type=int, default=100, help='Max results')
    parser.add_argument('--pretty', '-p', action='store_true', help='Pretty print JSON output')
    parser.add_argument('--raw', action='store_true', help='Output raw Loki response (not just log lines)')
    
    args = parser.parse_args()
    
    # Parse Grafana URL if provided
    if args.url:
        query, time_range = parse_grafana_url(args.url)
        if not query:
            print("❌ Failed to parse Grafana URL", file=sys.stderr)
            sys.exit(1)
        # Override time if not explicitly set
        if not args.time:
            args.time = time_range
        print(f"✓ Extracted from Grafana URL:", file=sys.stderr)
        print(f"   Query: {query}", file=sys.stderr)
        print(f"   Time: {args.time}", file=sys.stderr)
    # Build or use query
    elif args.query:
        query = args.query
        if not args.time:
            args.time = '1h'
    elif args.app:
        query = f'{{app="{args.app}"}}'
        if args.filter:
            query += f' |= `{args.filter}`'
        if args.severity:
            query += f' | json | severity="{args.severity}"'
        if not args.time:
            args.time = '1h'
    else:
        parser.print_help()
        sys.exit(1)
    
    # Show query info (if not already shown from URL parsing)
    if not args.url:
        print(f"🔍 Query Info:", file=sys.stderr)
        print(f"   Query: {query}", file=sys.stderr)
        print(f"   Range: Last {args.time}", file=sys.stderr)
    
    print(f"   Limit: {args.limit}", file=sys.stderr)
    print(file=sys.stderr)
    print(f"🔍 Querying Loki...", file=sys.stderr)
    
    # Query Loki
    data = query_loki_raw(query, time_range=args.time, limit=args.limit)
    
    # Output results
    if args.raw:
        # Output complete Loki response
        print(json.dumps(data, indent=2 if args.pretty else None))
    else:
        # Extract and output just the log lines
        logs = []
        for stream in data.get('data', {}).get('result', []):
            for value in stream.get('values', []):
                timestamp, log_line = value
                logs.append((timestamp, log_line))
        
        if not logs:
            print("No logs found", file=sys.stderr)
            sys.exit(0)
        
        print(f"✓ Found {len(logs)} log(s)\n", file=sys.stderr)
        
        # Output logs
        for timestamp, log_line in logs:
            try:
                # Parse and re-output JSON (validates it's JSON)
                log_obj = json.loads(log_line)
                if args.pretty:
                    print(json.dumps(log_obj, indent=2))
                else:
                    print(json.dumps(log_obj))
            except json.JSONDecodeError:
                # Not JSON, output as-is
                print(log_line)

if __name__ == '__main__':
    main()
