import nginx_pb2_grpc
 
 
-def run(args):
+def run(args, grpc_port='50054'):
     # get pod ip for grpc
     pod_ip = get_podip(args['service_name'])
     if pod_ip == '':
-        return "Cant find service with name: {}".format(args['service_name'])
-    nginx_grpc = pod_ip + ':50054'
+        return "Cant find service: {}".format(args['service_name'])
+    nginx_grpc = pod_ip + ':' + grpc_port
     channel = grpc.insecure_channel(nginx_grpc)
     stub = nginx_pb2_grpc.ControllerStub(channel)
 
     # modify config
-    if args['service_type'] == 'lb':
-        modify_lb(stub)
+    if args['service_type'] == 'lbv1':
+        slb_list = pickle.dumps(
+                    ['clover-server1:9180', 'clover-server2:9180'])
+        modify_lb(stub, slb_list)
+    if args['service_type'] == 'lbv2':
+        slb_list = pickle.dumps(
+                    ['clover-server4:9180', 'clover-server5:9180'])
+        modify_lb(stub, slb_list)
     elif args['service_type'] == 'proxy':
         modify_proxy(stub)
     elif args['service_type'] == 'server':
 
 
 def get_podip(pod_name):
-    config.load_kube_config()
-    v1 = client.CoreV1Api()
-    ret = v1.list_pod_for_all_namespaces(watch=False)
     ip = ''
-    for i in ret.items:
-        if i.metadata.name.lower().find(pod_name.lower()) != -1:
-            print(i.status.pod_ip)
-            ip = i.status.pod_ip
+    if pod_name != '':
+        config.load_kube_config()
+        v1 = client.CoreV1Api()
+        ret = v1.list_pod_for_all_namespaces(watch=False)
+        for i in ret.items:
+            if i.metadata.name.lower().find(pod_name.lower()) != -1:
+                print("Pod IP: {}".format(i.status.pod_ip))
+                ip = i.status.pod_ip
+                return str(ip)
     return str(ip)
 
 
     print(response.message)
 
 
-def modify_lb(stub):
-    slb_list = pickle.dumps(
-                    ['clover-server1', 'clover-server2', 'clover-server3'])
+def modify_lb(stub, slb_list):
     response = stub.ModifyLB(nginx_pb2.ConfigLB(
-            server_port='9188', server_name='http-lb',
+            server_port='9180', server_name='http-lb',
             slb_list=slb_list,
             slb_group='cloverlb', lb_path='/'))
     print(response.message)
 
             # self.template_file = 'templates/proxy.template'
             self.template_file = '/grpc/templates/proxy.template'
             self.ModifyProxy(nginx_pb2.ConfigProxy(
-                server_port='9180', server_name='http-proxy',
-                location_path='/', proxy_path='http://clover-server:9180',
+                server_port='9180', server_name='proxy-access-control',
+                location_path='/', proxy_path='http://http-lb:9180',
                 mirror_path='http://snort-ids:80'), "")
         if service_type == "server":
             # self.template_file = 'templates/server.template'
             # self.template_file = 'templates/lb.template'
             self.template_file = '/grpc/templates/lb.template'
             slb_list = pickle.dumps(
-                    ['clover-server1', 'clover-server2', 'clover-server3'])
+                    ['clover-server1:9180', 'clover-server2:9180',
+                        'clover-server3:9180'])
             self.ModifyLB(nginx_pb2.ConfigLB(
                 server_port='9180', server_name='http-lb',
                 slb_list=slb_list,
 
 
     upstream {{ slb_group }} {
         {%- for item in slb_list %}
-            server {{ item }};
+            server 127.0.0.1:800{{ loop.index }};
         {%- endfor %}
     }
 
+    {%- for item in slb_list %}
+    server {
+        listen 800{{ loop.index }};
+        location / {
+           proxy_http_version 1.1;
+           proxy_pass http://{{ item }};
+        }
+    }
+    {%- endfor %}
+
     server {
         listen {{ server_port }};
         server_name {{ server_name }};
             proxy_http_version 1.1;
             proxy_set_header Upgrade $http_upgrade;
             proxy_set_header Connection "upgrade";
+            proxy_set_header X-Forwarded-For $remote_addr;
+            proxy_set_header X-Forwarded-Proto http;
+            proxy_set_header X-Real-IP $remote_addr;
             proxy_pass http://{{ slb_group }};
         }
 
 
             proxy_http_version 1.1;
             proxy_set_header Upgrade $http_upgrade;
             proxy_set_header Connection "upgrade";
+            proxy_set_header X-Forwarded-For $remote_addr;
+            proxy_set_header X-Forwarded-Proto http;
+            proxy_set_header X-Real-IP $remote_addr;
             proxy_pass {{ proxy_path }};
 
             post_action @post_ids;
             proxy_http_version 1.1;
             proxy_set_header Upgrade $http_upgrade;
             proxy_set_header Connection "upgrade";
+            proxy_set_header X-Forwarded-For $remote_addr;
+            proxy_set_header X-Forwarded-Proto http;
+            proxy_set_header X-Real-IP $remote_addr;
             proxy_pass {{ mirror_path }};
         }
 
 
     grpc_port = '50054'
     if args['service_type'] == 'lb':
         out_file = 'lb.yaml'
-        deploy_name = 'http-lb'
+        if args['deploy_name'] == 'default':
+            args['deploy_name'] = 'http-lb'
     elif args['service_type'] == 'proxy':
         out_file = 'proxy.yaml'
-        deploy_name = 'proxy-access-control'
+        if args['deploy_name'] == 'default':
+            args['deploy_name'] = 'proxy-access-control'
     elif args['service_type'] == 'server':
         out_file = 'server.yaml'
-        deploy_name = 'clover-server'
+        if args['deploy_name'] == 'default':
+            args['deploy_name'] = 'clover-server'
     else:
         return "Invalid service type: {}".format(args['service_type'])
 
             image_path=args['image_path'],
             image_name=args['image_name'],
             image_tag=args['image_tag'],
-            deploy_name=deploy_name,
+            deploy_name=args['deploy_name'],
             server_port=server_port,
             grpc_port=grpc_port
         )
     parser.add_argument(
             '--image_tag', default='latest',
             help='The image tag to use')
+    parser.add_argument(
+            '--deploy_name', default='default',
+            help='The deployment name to use')
     args = parser.parse_args()
     print(render_yaml(vars(args)))