import os
import configparser
from concurrent.futures import ThreadPoolExecutor
from flask import Flask, jsonify, render_template
from kubernetes import client, config
from kubernetes.client.rest import ApiException

def get_version():
    try:
        ini_path = os.path.join(os.path.dirname(__file__), 'version.ini')
        if os.path.exists(ini_path):
            parser = configparser.ConfigParser()
            parser.read(ini_path)
            return parser.get('metadata', 'version', fallback='0.5')
    except Exception as e:
        print(f"Error reading version.ini: {e}")
    return '0.5'

app = Flask(__name__)

_v1_client = None
_custom_client = None

def get_k8s_client():
    global _v1_client, _custom_client
    if _v1_client is not None and _custom_client is not None:
        return _v1_client, _custom_client
    try:
        config.load_incluster_config()
    except config.ConfigException:
        try:
            config.load_kube_config()
        except Exception as e:
            print(f"Error loading kubeconfig: {e}")
            loaded = False
            for fallback_path in ["/home/nonroot/.kube/config", "/root/.kube/config"]:
                try:
                    if os.path.exists(fallback_path):
                        config.load_kube_config(config_file=fallback_path)
                        loaded = True
                        break
                except Exception:
                    pass
            if not loaded:
                raise RuntimeError(f"Failed to load kubeconfig. Make sure the file exists and is accessible. Error: {e}")
    _v1_client = client.CoreV1Api()
    _custom_client = client.CustomObjectsApi()
    return _v1_client, _custom_client

def parse_cpu(cpu_str):
    if not cpu_str: return 0
    cpu_str = str(cpu_str)
    if cpu_str.endswith('m'):
        return int(cpu_str[:-1])
    if cpu_str.endswith('n'):
        return int(cpu_str[:-1]) // 1000000
    try:
        return int(float(cpu_str) * 1000)
    except:
        return 0

def parse_memory(mem_str):
    if not mem_str: return 0
    mem_str = str(mem_str)
    if mem_str.endswith('Ki'):
        return int(mem_str[:-2]) * 1024
    if mem_str.endswith('Mi'):
        return int(mem_str[:-2]) * 1024**2
    if mem_str.endswith('Gi'):
        return int(mem_str[:-2]) * 1024**3
    try:
        return int(mem_str)
    except:
        return 0

K8S_TIMEOUT = 5.0
METRICS_TIMEOUT = 3.0

def fetch_nodes(v1):
    return v1.list_node(_request_timeout=K8S_TIMEOUT).items

def fetch_pods(v1):
    return v1.list_pod_for_all_namespaces(_request_timeout=K8S_TIMEOUT).items

def fetch_node_metrics(custom):
    node_metrics = {}
    try:
        nm = custom.list_cluster_custom_object("metrics.k8s.io", "v1beta1", "nodes", _request_timeout=METRICS_TIMEOUT)
        for item in nm.get('items', []):
            node_metrics[item['metadata']['name']] = item['usage']
    except Exception as e:
        print(f"Metrics server node fetch error or unavailable: {e}")
    return node_metrics

def fetch_pod_metrics(custom):
    pod_metrics = {}
    try:
        pm = custom.list_cluster_custom_object("metrics.k8s.io", "v1beta1", "pods", _request_timeout=METRICS_TIMEOUT)
        for item in pm.get('items', []):
            key = f"{item['metadata']['namespace']}/{item['metadata']['name']}"
            pod_metrics[key] = item['containers']
    except Exception as e:
        print(f"Metrics server pod fetch error or unavailable: {e}")
    return pod_metrics

@app.route('/api/cluster-data')
def cluster_data():
    try:
        v1, custom = get_k8s_client()

        with ThreadPoolExecutor(max_workers=4) as executor:
            fut_nodes = executor.submit(fetch_nodes, v1)
            fut_pods = executor.submit(fetch_pods, v1)
            fut_node_metrics = executor.submit(fetch_node_metrics, custom)
            fut_pod_metrics = executor.submit(fetch_pod_metrics, custom)

            nodes = fut_nodes.result()
            pods = fut_pods.result()
            node_metrics = fut_node_metrics.result()
            pod_metrics = fut_pod_metrics.result()

        cluster = {
            "name": "Cluster",
            "children": []
        }

        pods_by_node = {}
        for pod in pods:
            nname = pod.spec.node_name
            if nname:
                pods_by_node.setdefault(nname, []).append(pod)
            
        for node in nodes:
            node_name = node.metadata.name

            internal_ip = "N/A"
            if node.status and getattr(node.status, 'addresses', None):
                for addr in node.status.addresses:
                    addr_type = getattr(addr, 'type', None) or (addr.get('type') if isinstance(addr, dict) else None)
                    addr_val = getattr(addr, 'address', None) or (addr.get('address') if isinstance(addr, dict) else None)
                    if addr_type == 'InternalIP':
                        internal_ip = addr_val
                        break

            os_image = "N/A"
            if node.status and getattr(node.status, 'node_info', None):
                node_info = node.status.node_info
                if isinstance(node_info, dict):
                    os_image = node_info.get('osImage') or node_info.get('os_image') or "N/A"
                else:
                    os_image = getattr(node_info, 'os_image', None) or getattr(node_info, 'os_image', "N/A")

            labels = node.metadata.labels or {}
            roles = []
            for k in labels:
                if k.startswith("node-role.kubernetes.io/"):
                    role_name = k.split("/", 1)[1]
                    if role_name:
                        roles.append(role_name)

            is_control_plane = any(r in ["control-plane", "master"] for r in roles) or ("control-plane" in node_name) or ("master" in node_name)
            if is_control_plane and not roles:
                roles = ["control-plane" if "control-plane" in node_name else "master"]
            elif not roles:
                roles = ["worker"]

            cpu_allocatable = parse_cpu(node.status.allocatable.get('cpu', '0'))
            mem_allocatable = parse_memory(node.status.allocatable.get('memory', '0'))

            node_usage_cpu = parse_cpu(node_metrics.get(node_name, {}).get('cpu', '0'))
            node_usage_mem = parse_memory(node_metrics.get(node_name, {}).get('memory', '0'))

            node_pods = pods_by_node.get(node_name, [])
            non_terminated_pods = [
                p for p in node_pods
                if p.status and getattr(p.status, 'phase', None) not in ['Succeeded', 'Failed']
            ]
            non_terminated_pod_count = len(non_terminated_pods)

            node_data = {
                "name": node_name,
                "type": "node",
                "roles": roles,
                "is_control_plane": is_control_plane,
                "isControlPlane": is_control_plane,
                "internal_ip": internal_ip,
                "internalIp": internal_ip,
                "os_image": os_image,
                "osImage": os_image,
                "non_terminated_pods": non_terminated_pod_count,
                "nonTerminatedPods": non_terminated_pod_count,
                "pod_count": non_terminated_pod_count,
                "podCount": non_terminated_pod_count,
                "cpu_allocatable": cpu_allocatable,
                "mem_allocatable": mem_allocatable,
                "cpu_usage": node_usage_cpu,
                "mem_usage": node_usage_mem,
                "children": []
            }
            
            node_cpu_total = 0
            node_mem_total = 0
            
            for pod in node_pods:
                pod_cpu_req = 0
                pod_mem_req = 0
                
                if pod.spec.containers:
                    for c in pod.spec.containers:
                        if c.resources and c.resources.requests:
                            pod_cpu_req += parse_cpu(c.resources.requests.get('cpu', '0'))
                            pod_mem_req += parse_memory(c.resources.requests.get('memory', '0'))
                            
                key = f"{pod.metadata.namespace}/{pod.metadata.name}"
                if key in pod_metrics:
                    pod_cpu = sum([parse_cpu(c['usage'].get('cpu', '0')) for c in pod_metrics[key]])
                    pod_mem = sum([parse_memory(c['usage'].get('memory', '0')) for c in pod_metrics[key]])
                else:
                    pod_cpu = pod_cpu_req
                    pod_mem = pod_mem_req
                    
                # D3 Treemap doesn't like 0 values, so minimum size 1
                pod_cpu_val = max(pod_cpu, 1)
                pod_mem_val = max(pod_mem, 1)
                
                node_cpu_total += pod_cpu
                node_mem_total += pod_mem
                    
                node_data["children"].append({
                    "name": pod.metadata.name,
                    "namespace": pod.metadata.namespace,
                    "type": "pod",
                    "cpu": pod_cpu_val,
                    "memory": pod_mem_val,
                    "cpu_raw": pod_cpu,
                    "memory_raw": pod_mem,
                    "status": pod.status.phase
                })
            
            # Optionally add an "Unused" block so the node size matches allocatable
            unused_cpu = max(0, cpu_allocatable - node_cpu_total)
            unused_mem = max(0, mem_allocatable - node_mem_total)
            
            if unused_cpu > 0 or unused_mem > 0:
                 node_data["children"].append({
                     "name": "Available Capacity",
                     "namespace": "",
                     "type": "unused",
                     "cpu": max(unused_cpu, 1),
                     "memory": max(unused_mem, 1),
                     "cpu_raw": unused_cpu,
                     "memory_raw": unused_mem,
                     "status": "Available"
                 })
                 
            cluster["children"].append(node_data)
            
        return jsonify(cluster)
        
    except Exception as e:
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500

@app.route('/api/version')
def version_route():
    return jsonify({"version": get_version()})

@app.route('/healthz')
def healthz():
    return jsonify({"status": "ok"}), 200

@app.route('/')
def index():
    return render_template('index.html', version=get_version())

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000)
