#!/usr/bin/env python3
"""
Generate TinyRDM connections.yaml from sdm status --json output
Organizes Redis connections into Internal, Stage, and Production groups
Excludes activate and readonly connections
"""

import subprocess
import json
import re
import os
from collections import defaultdict

def run_sdm_status():
    try:
        result = subprocess.run(['sdm', 'status', '--json'], capture_output=True, text=True)
        if result.returncode != 0:
            print(f"Error running sdm status: {result.stderr}")
            return None
        return json.loads(result.stdout)
    except FileNotFoundError:
        print("Error: sdm command not found")
        return None
    except json.JSONDecodeError as e:
        print(f"Error parsing JSON: {e}")
        return None

def parse_redis_connections(sdm_data):
    connections = []
    
    for item in sdm_data:
        if item.get('type') == 'redis':
            name = item.get('name', '')
            address = item.get('address', '')
            
            if ('activate-cache' in name or 'activate-readonly' in name or 
                'readonly-redis' in name):
                continue
            
            if ':' in address:
                addr, port = address.split(':')
                port = int(port)
            else:
                continue
            
            env_info = parse_connection_name(name)
            if env_info:
                connections.append({
                    'name': name,
                    'addr': addr,
                    'port': port,
                    'environment': env_info['environment'],
                    'customer': env_info['customer'],
                    'stage': env_info['stage']
                })
    
    return connections

def parse_connection_name(name):
    pattern = r'oc-([^-]+)-([^-]+)-.*'
    match = re.match(pattern, name)
    
    if not match:
        return None
    
    environment = match.group(1)
    customer = match.group(2)
    
    if environment in ['dev', 'nextrc', 'nextrc2']:
        stage = 'internal'
    elif environment in ['stage', 'stage2', 'uat']:
        stage = 'stage'
    elif environment in ['prod', 'prod2']:
        stage = 'production'
    else:
        stage = 'unknown'
    
    return {
        'environment': environment,
        'customer': customer,
        'stage': stage
    }

def create_connection_yaml(name, addr, port, stage):
    colors = {
        'internal': '#4ECF60',
        'stage': '#FFA500',
        'production': '#FF0000'
    }
    
    color = colors.get(stage, '#808080')
    
    return f"""    - name: {name}
      last_db: 0
      network: tcp
      addr: {addr}
      port: {port}
      default_filter: '*'
      key_separator: ':'
      conn_timeout: 60
      exec_timeout: 60
      db_filter_type: none
      load_size: 10000
      mark_color: '{color}'
      refresh_interval: 5"""

def group_connections(connections):
    groups = defaultdict(list)
    
    for conn in connections:
        # Mirror DBeaver grouping: internal/features -> 'internal', otherwise by customer, else 'other'
        if conn['customer'] == 'internal' or conn['customer'].startswith('feature'):
            group_name = 'internal'
        elif conn['stage'] in ['stage', 'production']:
            group_name = conn['customer']
        else:
            group_name = 'other'
        
        # Keep stage-based color coding
        if conn['stage'] == 'internal':
            conn_stage = 'internal'
        elif conn['stage'] == 'stage':
            conn_stage = 'stage'
        elif conn['stage'] == 'production':
            conn_stage = 'production'
        else:
            conn_stage = 'unknown'
        
        conn_yaml = create_connection_yaml(
            conn['name'],
            conn['addr'],
            conn['port'],
            conn_stage
        )
        
        groups[group_name].append(conn_yaml)
    
    return groups

def generate_yaml_content(groups):
    yaml_lines = []
    # Prefer 'internal' first, then alphabetical by customer
    ordered_groups = []
    if 'internal' in groups:
        ordered_groups.append('internal')
    ordered_groups.extend(sorted([g for g in groups.keys() if g != 'internal']))
    
    for group_name in ordered_groups:
        yaml_lines.append(f"- name: {group_name}")
        yaml_lines.append("  last_db: 0")
        yaml_lines.append("  type: group")
        yaml_lines.append("  connections:")
        
        for conn_yaml in groups[group_name]:
            yaml_lines.append(conn_yaml)
        
        yaml_lines.append("")
    
    return '\n'.join(yaml_lines)

def main():
    print("Generating TinyRDM connections.yaml from sdm status --json...")
    
    sdm_data = run_sdm_status()
    if not sdm_data:
        return
    
    connections = parse_redis_connections(sdm_data)
    print(f"Found {len(connections)} Redis connections")
    
    groups = group_connections(connections)
    
    for group_name, conns in groups.items():
        print(f"  {group_name}: {len(conns)} connections")
    
    yaml_content = generate_yaml_content(groups)
    
    output_file = os.path.expanduser('~/.config/TinyRDM/connections.yaml')
    try:
        with open(output_file, 'w') as f:
            f.write(yaml_content)
        print(f"✅ Successfully generated {output_file}")
    except Exception as e:
        print(f"❌ Error writing file: {e}")

if __name__ == '__main__':
    main()