import yaml

def str_presenter(dumper, data):
    if '\n' in data:
        return dumper.represent_scalar('tag:yaml.org,2002:str', data, style='|')
    return dumper.represent_scalar('tag:yaml.org,2002:str', data)

yaml.add_representer(str, str_presenter)
yaml.representer.SafeRepresenter.add_representer(str, str_presenter)

with open("homelab-gitops-repo/manifests/aerodrome/prometheus/prometheus.yaml", "r", encoding="utf-8") as f:
    docs = list(yaml.safe_load_all(f.read()))

# Doc 0: PVC
# Doc 1: ConfigMap
cm = docs[1]
prom_cfg = yaml.safe_load(cm['data']['prometheus.yml'])

scrape_configs = prom_cfg.get('scrape_configs', [])

# Check if node-exporter already exists
if not any(j.get('job_name') == 'node-exporter' for j in scrape_configs):
    scrape_configs.append({
        'job_name': 'node-exporter',
        'metrics_path': '/metrics',
        'static_configs': [
            {
                'targets': [
                    '140.44.4.11:9100', '140.44.4.12:9100', '140.44.4.13:9100',
                    '140.44.4.21:9100', '140.44.4.22:9100', '140.44.4.23:9100',
                    '140.44.4.24:9100', '140.44.4.25:9100', '140.44.4.26:9100',
                    '140.44.4.27:9100'
                ],
                'labels': {
                    'cluster': 'aerodrome'
                }
            },
            {
                'targets': [
                    '140.44.4.41:9100', '140.44.4.42:9100', '140.44.4.43:9100',
                    '140.44.4.51:9100', '140.44.4.52:9100', '140.44.4.53:9100',
                    '140.44.4.54:9100', '140.44.4.55:9100', '140.44.4.56:9100',
                    '140.44.4.57:9100'
                ],
                'labels': {
                    'cluster': 'spaceport'
                }
            }
        ]
    })

# Check if kubernetes-cadvisor already exists
if not any(j.get('job_name') == 'kubernetes-cadvisor' for j in scrape_configs):
    scrape_configs.append({
        'job_name': 'kubernetes-cadvisor',
        'scheme': 'https',
        'tls_config': {
            'ca_file': '/var/run/secrets/kubernetes.io/serviceaccount/ca.crt'
        },
        'bearer_token_file': '/var/run/secrets/kubernetes.io/serviceaccount/token',
        'kubernetes_sd_configs': [
            {'role': 'node'}
        ],
        'relabel_configs': [
            {
                'action': 'labelmap',
                'regex': '__meta_kubernetes_node_label_(.+)'
            },
            {
                'target_label': '__address__',
                'replacement': 'kubernetes.default.svc:443'
            },
            {
                'source_labels': ['__meta_kubernetes_node_name'],
                'regex': '(.+)',
                'target_label': '__metrics_path__',
                'replacement': '/api/v1/nodes/${1}/proxy/metrics/cadvisor'
            }
        ]
    })

prom_cfg['scrape_configs'] = scrape_configs
cm['data']['prometheus.yml'] = yaml.dump(prom_cfg, sort_keys=False, default_flow_style=False)

# Doc 2: Deployment
deploy = docs[2]
deploy['spec']['template']['spec']['serviceAccountName'] = 'prometheus'

# Add SA, ClusterRole, ClusterRoleBinding if not in docs
sa = {
    'apiVersion': 'v1',
    'kind': 'ServiceAccount',
    'metadata': {
        'name': 'prometheus',
        'namespace': 'default',
        'labels': {'app.kubernetes.io/name': 'prometheus'}
    }
}

cr = {
    'apiVersion': 'rbac.authorization.k8s.io/v1',
    'kind': 'ClusterRole',
    'metadata': {
        'name': 'prometheus',
        'labels': {'app.kubernetes.io/name': 'prometheus'}
    },
    'rules': [
        {
            'apiGroups': [''],
            'resources': ['nodes', 'nodes/proxy', 'nodes/metrics', 'services', 'endpoints', 'pods'],
            'verbs': ['get', 'list', 'watch']
        },
        {
            'nonResourceURLs': ['/metrics', '/metrics/cadvisor'],
            'verbs': ['get']
        }
    ]
}

crb = {
    'apiVersion': 'rbac.authorization.k8s.io/v1',
    'kind': 'ClusterRoleBinding',
    'metadata': {
        'name': 'prometheus',
        'labels': {'app.kubernetes.io/name': 'prometheus'}
    },
    'roleRef': {
        'apiGroup': 'rbac.authorization.k8s.io',
        'kind': 'ClusterRole',
        'name': 'prometheus'
    },
    'subjects': [
        {
            'kind': 'ServiceAccount',
            'name': 'prometheus',
            'namespace': 'default'
        }
    ]
}

new_docs = [sa, cr, crb] + docs

with open("homelab-gitops-repo/manifests/aerodrome/prometheus/prometheus.yaml", "w", encoding="utf-8") as f:
    yaml.dump_all(new_docs, f, sort_keys=False, default_flow_style=False)

print("Updated manifests/aerodrome/prometheus/prometheus.yaml successfully!")