poc1013

def run_mc_command(alias: str, insecure: bool = False) -> Dict:
    """Run mc admin info command and return JSON output"""
    cmd = f"mc admin info {alias} --json"
    if insecure:
        cmd += " --insecure"
    
    result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
    if result.returncode != 0:
        print(f"Error running mc command: {result.stderr}", file=sys.stderr)
        sys.exit(1)
    
    try:
        return json.loads(result.stdout)
    except json.JSONDecodeError as e:
        print(f"Error parsing mc output as JSON: {e}", file=sys.stderr)
        print(f"Output was: {result.stdout}", file=sys.stderr)
        sys.exit(1)


def get_pod_node_mapping(namespace: str, objectstore_name: str) -> Dict[str, str]:
    """
    Get mapping of pod names to node names using kubectl
    Returns dict mapping pod name to node name
    """
    pod_to_node = {}
    
    # Try different label selectors
    label_selectors = [
        f'v1.min.io/tenant={objectstore_name}',  # MinIO Operator label
        f'app=minio',  # Generic MinIO label
        f'app.kubernetes.io/name=minio',  # Standard k8s label
        f'app.kubernetes.io/instance={objectstore_name}',  # Instance label
    ]
    
    for selector in label_selectors:
        try:
            result = subprocess.run(
                ['kubectl', '-n', namespace, 'get', 'pods', '-l', selector, '-o', 'json'],
                capture_output=True,
                text=True,
                check=True
            )
            
            pods_info = json.loads(result.stdout)
            items = pods_info.get('items', [])
            
            if items:
                print(f"  Found {len(items)} pods with selector '{selector}'")
                for pod in items:
                    pod_name = pod['metadata']['name']
                    node_name = pod['spec'].get('nodeName', 'unknown')
                    # Only add if it looks like a MinIO pod (has 'minio' in name or matches objectstore name)
                    if objectstore_name in pod_name or 'minio' in pod_name.lower():
                        pod_to_node[pod_name] = node_name
                
                if pod_to_node:
                    break  # Found pods, stop trying other selectors
                    
        except subprocess.CalledProcessError:
            continue  # Try next selector
        except json.JSONDecodeError:
            continue
    
    # If no labeled pods found, try to get all pods and filter by name
    if not pod_to_node:
        try:
            print(f"  No pods found with standard labels, trying name-based search...")
            result = subprocess.run(
                ['kubectl', '-n', namespace, 'get', 'pods', '-o', 'json'],
                capture_output=True,
                text=True,
                check=True
            )
            
            pods_info = json.loads(result.stdout)
            for pod in pods_info.get('items', []):
                pod_name = pod['metadata']['name']
                # Filter for MinIO-like pod names
                if (objectstore_name in pod_name or 
                    'minio' in pod_name.lower() or 
                    pod_name.startswith(f'{objectstore_name}-pool')):
                    node_name = pod['spec'].get('nodeName', 'unknown')
                    pod_to_node[pod_name] = node_name
            
            if pod_to_node:
                print(f"  Found {len(pod_to_node)} MinIO pods by name pattern")
                
        except (subprocess.CalledProcessError, json.JSONDecodeError) as e:
            print(f"Warning: Could not get pod information: {e}", file=sys.stderr)
    
    return pod_to_node


def analyze_erasure_sets(mc_info: Dict, nodes_by_label: Dict[str, List[str]], 
                         node_label_key: str, namespace: str = None, 
                         objectstore_name: str = None) -> Dict:
    """Analyze erasure set distribution from mc admin info output"""
    
    # Extract servers and drives information
    servers = mc_info.get('info', {}).get('servers', [])
    
    if not servers:
        print("Error: No servers found in mc admin info output", file=sys.stderr)
        return {}
    
    # Build erasure set mapping
    erasure_sets = defaultdict(list)  # erasure_set_id -> list of drive info
    node_to_rack = {}  # node -> rack mapping
    
    # Build node to rack mapping
    for rack, nodes in nodes_by_label.items():
        for node in nodes:
            node_to_rack[node] = rack
    
    # Get pod to node mapping if namespace and objectstore are provided
    pod_to_node = {}
    if namespace and objectstore_name:
        pod_to_node = get_pod_node_mapping(namespace, objectstore_name)
        print(f"Found {len(pod_to_node)} pods in namespace {namespace}")
        for pod, node in pod_to_node.items():
            print(f"  Pod {pod} is running on node {node}")
    
    print(f"\nProcessing {len(servers)} servers from mc admin info...")
    
    # Process each server
    for i, server in enumerate(servers):
        # Extract endpoint (e.g., "minio-pool-0-0.minio-hl.minio-tenant.svc.cluster.local:9000")
        endpoint = server.get('endpoint', '')
        
        # Debug: show raw endpoint
        if i < 3 or 'unknown' in endpoint:  # Show first few or problematic ones
            print(f"  Server {i}: endpoint = '{endpoint}'")
        
        # Try to extract pod name from endpoint
        # Common patterns:
        # 1. "pod-name.service.namespace.svc.cluster.local:port"
        # 2. "pod-name:port"
        # 3. "node-name:port" (direct node access)
        # 4. "myminio-pool-0-0" (just pod name)
        
        # Remove port if present
        pod_or_node_name = endpoint.split(':')[0] if ':' in endpoint else endpoint
        # Remove domain parts if present
        pod_or_node_name = pod_or_node_name.split('.')[0] if '.' in pod_or_node_name else pod_or_node_name
        
        # Determine the actual node
        actual_node = 'unknown'
        
        # First try exact match with pod names
        if pod_or_node_name in pod_to_node:
            actual_node = pod_to_node[pod_or_node_name]
            if i < 3:  # Show first few for debugging
                print(f"    -> Matched pod '{pod_or_node_name}' -> Node {actual_node}")
        # Then try node names
        elif pod_or_node_name in node_to_rack:
            actual_node = pod_or_node_name
            if i < 3:
                print(f"    -> Direct node {actual_node}")
        else:
            # Try fuzzy matching with pod names
            # Handle cases where endpoint might have different naming
            matched = False
            for pod_name, node_name in pod_to_node.items():
                # Try various matching strategies
                if (pod_name == pod_or_node_name or 
                    pod_or_node_name in pod_name or
                    pod_name in pod_or_node_name or
                    # Handle pool naming differences (e.g., myminio-pool-0-0 vs minio-pool-0-0)
                    pod_name.split('-pool-')[-1] == pod_or_node_name.split('-pool-')[-1] if '-pool-' in pod_name and '-pool-' in pod_or_node_name else False):
                    actual_node = node_name
                    if i < 3:
                        print(f"    -> Fuzzy matched to pod '{pod_name}' -> Node {actual_node}")
                    matched = True
                    break
            
            if not matched:
                print(f"  WARNING: Could not determine node for endpoint '{endpoint}'")
                print(f"    Extracted name: '{pod_or_node_name}'")
                if pod_to_node:
                    print(f"    Available pods: {list(pod_to_node.keys())[:3]}...")
                if node_to_rack:
                    print(f"    Available nodes: {list(node_to_rack.keys())[:3]}...")
        
        # Get rack for this node
        rack = node_to_rack.get(actual_node, 'unknown')
        
        # Process drives
        drives = server.get('drives', [])
        for drive in drives:
            # Look for set_index field (this is what MinIO uses in newer versions)
            erasure_set_id = drive.get('set_index')
            if erasure_set_id is None:
                # Fallback to 'set' field if set_index not found
                erasure_set_id = drive.get('set')
            
            if erasure_set_id is not None:
                erasure_sets[erasure_set_id].append({
                    'node': actual_node,
                    'rack': rack,
                    'drive': drive.get('path', ''),
                    'state': drive.get('state', 'unknown'),
                    'endpoint': endpoint,
                    'pod_or_node': pod_or_node_name
                })
    
    print(f"\nFound {len(erasure_sets)} erasure sets with {sum(len(drives) for drives in erasure_sets.values())} total drives")
    
    return erasure_sets


def generate_audit_report(erasure_sets: Dict, nodes_by_label: Dict[str, List[str]], 
                          objectstore_name: str) -> str:
    """Generate a text-based audit report of erasure set placement"""
    
    report = []
    report.append("=" * 80)
    report.append(f"MinIO Erasure Set Placement Audit Report")
    report.append(f"ObjectStore: {objectstore_name}")
    report.append("=" * 80)
    report.append("")
    
    # Summary statistics
    total_sets = len(erasure_sets)
    total_drives = sum(len(drives) for drives in erasure_sets.values())
    total_racks = len(nodes_by_label)
    total_nodes = sum(len(nodes) for nodes in nodes_by_label.values())
    
    report.append("📊 Summary Statistics:")
    report.append(f"  • Total Erasure Sets: {total_sets}")
    report.append(f"  • Total Drives: {total_drives}")
    report.append(f"  • Total Racks: {total_racks}")
    report.append(f"  • Total Nodes: {total_nodes}")
    report.append("")
    
    # Analyze each erasure set
    report.append("🔍 Erasure Set Analysis:")
    report.append("-" * 40)
    
    issues_found = []
    
    for set_id in sorted(erasure_sets.keys()):
        drives = erasure_sets[set_id]
        report.append(f"\nErasure Set {set_id}:")
        report.append(f"  Drives: {len(drives)}")
        
        # Count drives per rack
        rack_counts = Counter(drive['rack'] for drive in drives)
        report.append(f"  Distribution across racks:")
        
        for rack in sorted(rack_counts.keys(), key=natural_sort_key):
            count = rack_counts[rack]
            report.append(f"    Rack {rack}: {count} drives")
            
            # Check for issues: multiple drives from same erasure set on same rack
            if count > 1:
                issue = f"⚠️  WARNING: Erasure set {set_id} has {count} drives on rack {rack}"
                issues_found.append(issue)
                report.append(f"    {issue}")
        
        # Show detailed drive placement
        report.append("  Drive details:")
        for drive in sorted(drives, key=lambda x: (x['rack'], x['node'], x['drive'])):
            state_icon = "✅" if drive['state'] == 'online' else "❌"
            report.append(f"    {state_icon} Rack {drive['rack']}, Node {drive['node']}: {drive['drive']}")
    
    # Issues summary
    report.append("")
    report.append("=" * 80)
    if issues_found:
        report.append("⚠️  ISSUES FOUND:")
        report.append("-" * 40)
        for issue in issues_found:
            report.append(issue)
        report.append("")
        report.append("RECOMMENDATION: Multiple drives from the same erasure set on the same rack")
        report.append("reduces fault tolerance. Consider redistributing pods across racks.")
    else:
        report.append("✅ No placement issues found!")
        report.append("All erasure sets are properly distributed across racks.")
    
    report.append("=" * 80)
    
    # Rack utilization
    report.append("")
    report.append("📈 Rack Utilization:")
    report.append("-" * 40)
    
    rack_drive_counts = defaultdict(int)
    for drives in erasure_sets.values():
        for drive in drives:
            rack_drive_counts[drive['rack']] += 1
    
    for rack in sorted(rack_drive_counts.keys(), key=natural_sort_key):
        count = rack_drive_counts[rack]
        bar = "█" * (count // 2) + ("▌" if count % 2 else "")
        report.append(f"  Rack {rack}: {bar} ({count} drives)")
    
    return "\n".join(report)


def audit_mode(args):
    """Run in audit mode to analyze existing MinIO deployment"""
    
    print("🔍 Running MinIO Placement Audit...")
    
    # Get MinIO server info
    print(f"Fetching MinIO server info from '{args.mc_alias}'...")
    mc_info = run_mc_command(args.mc_alias, args.insecure)
    
    # Get ObjectStore CR
    print(f"Fetching ObjectStore CR...")
    objectstore = get_objectstore_cr(args.name, args.namespace, args.file)
    objectstore_name = objectstore['metadata']['name']
    namespace = objectstore['metadata']['namespace']
    
    # Get nodes grouped by label (pass audit_mode=True and support JSON file)
    print(f"Fetching nodes...")
    nodes_by_label = get_nodes_by_label(args.node_label, audit_mode=True, json_file=args.nodes_json)
    
    if not nodes_by_label:
        print(f"Error: No nodes found at all", file=sys.stderr)
        sys.exit(1)
    
    # Analyze erasure sets
    print("Analyzing erasure set distribution...")
    erasure_sets = analyze_erasure_sets(mc_info, nodes_by_label, args.node_label or "rack", 
                                       namespace, objectstore_name)
    
    # Generate report
    report = generate_audit_report(erasure_sets, nodes_by_label, objectstore_name)
    
    # Output report
    print("\n" + report)
    
    # Optionally save to file
    if args.audit_output:
        with open(args.audit_output, 'w') as f:
            f.write(report)
        print(f"\n📄 Report saved to: {args.audit_output}")


def main():
    parser = argparse.ArgumentParser(
        description='Generate PVCs for MinIO ObjectStore with balanced placement'
    )
    parser.add_argument('--name', help='Name of the ObjectStore CR')
    parser.add_argument('--namespace', '-n', help='Namespace of the ObjectStore CR')
    parser.add_argument('--file', '-f', help='Path to ObjectStore CR YAML file')
    parser.add_argument('--node-label', 
                       help='Node label for placement buckets (e.g., v1.min.io/rack)')
    parser.add_argument('--nodes-json',
                       help='JSON file containing node configuration (rack -> [nodes])')
    parser.add_argument('--output', '-o', default='pvcs.yaml',
                       help='Output YAML file path (default: pvcs.yaml)')
    parser.add_argument('--html-output', 
                       help='Generate HTML visualization of placement')
    parser.add_argument('--base-storage-class', 
                       help='Base storage class to use as template')
    parser.add_argument('--directpv-drives-label',
                       help='DirectPV drive label to filter drives (e.g., directpv.min.io/tier=fast)')
    parser.add_argument('--storage-class-per-node', action='store_true',
                       help='Create one storage class per node (default: one per rack)')
    
    # Audit mode arguments
    parser.add_argument('--audit', action='store_true',
                       help='Run in audit mode to analyze existing MinIO deployment')
    parser.add_argument('--mc-alias', 
                       help='MinIO mc alias for admin info command (required for audit mode)')
    parser.add_argument('--insecure', action='store_true',
                       help='Use --insecure flag with mc command')
    parser.add_argument('--audit-output',
                       help='Save audit report to file')
    
    args = parser.parse_args()
    
    # Check for audit mode
    if args.audit:
        if not args.mc_alias:
            print("Error: --mc-alias is required for audit mode", file=sys.stderr)
            sys.exit(1)
        audit_mode(args)
        return
    
    # Validate required arguments for normal mode
    if not args.node_label and not args.nodes_json:
        print("Error: Either --node-label or --nodes-json is required", file=sys.stderr)
        sys.exit(1)
    
    if args.node_label and args.nodes_json:
        print("Error: Cannot specify both --node-label and --nodes-json", file=sys.stderr)
        sys.exit(1)
    
    # Normal mode: Generate PVCs