import matplotlib.pyplot as plt
import numpy as np

from io import BytesIO
from filelock import FileLock
import subprocess
import logging
from io import BytesIO
from PIL import Image
from pycocotools import mask as mask_utils
import torch
import sys
import base64
import threading

import numpy as np
from flask import Flask, request , make_response , jsonify
from PIL import Image , ImageOps
from flask_caching import Cache
import json
import cv2

from flask import Flask, request,  jsonify
import sys
import base64
import threading
from PIL import Image
from io import BytesIO
import cv2

from fvcore.common.timer import Timer
from fvcore.transforms.transform_util import to_numpy
from pycocotools import mask as maskUtils
import numpy as np
import os
import csv
import re

from k8s_logger import weshop_log
import time
project_dir = os.environ["PROJECT_DIR"]
app_port_num = os.environ["APP_PORT_NUM"]
# project_dir = "F:\weshop"
# app_port_num = "30000"



sd_models_path = os.path.join( project_dir, "models" ,"Stable-diffusion")
sd_vae_path = os.path.join( project_dir, "models" ,"VAE")
sd_lora_path = os.path.join( project_dir, "models" ,"Lora")

python_exe = os.path.join(project_dir, "bin" ,"python", "python.exe")


gpu_lock_file =os.path.join(project_dir, "resources", "lock", "lock.lock")

gpu_service_pid = os.getpid()

gpu_pid_file = os.path.join(project_dir, "bin" ,"gpu_service.txt")

path = os.path.join(os.path.dirname(__file__), "output.csv")
data_dict = {}

with open(path, newline='') as csvfile:
    reader = csv.reader(csvfile)
    next(reader)
    for row in reader:
        if len(row) >= 2:
            text = row[0]
            count = int(row[1])
            data_dict[text] = count


with open(gpu_pid_file, "w") as file:
    file.write(str(gpu_service_pid))

def build_success_response(numpy):
   return jsonify({
      "success" : True ,
      "message" : "success",
      "data" : numpy
   })

def build_error_response(message):
   return jsonify({
      "success" : False ,
      "message" : message,
      "data" : message
   })



logger = logging.getLogger(__name__)

app = Flask(__name__)
app.logger.setLevel('INFO')
app.config['CACHE_TYPE'] = 'simple'
app.config['CACHE_DEFAULT_TIMEOUT'] = 3600 
cache = Cache(app)

# segment anything
# device = torch.device('cuda')
# segment_model = SegmentAnything(device, arch="vit_h", task = 'online', output_mode='uncompressed_rle')

# lock process
segment_pipe_lock = threading.Lock()
detection_pipe_lock = threading.Lock()
caption_pipe_lock = threading.Lock()
everything_pipe_lock = threading.Lock()

def base64_to_pil(info):
    decoded_bytes = base64.b64decode(info)
    image_stream = BytesIO(decoded_bytes)
    image = Image.open(image_stream).convert('RGB')
    return image



def show_anns(anns):
    if len(anns) == 0:
        return
    sorted_anns = sorted(anns, key=(lambda x: x['area']), reverse=True)
    ax = plt.gca()
    ax.set_autoscale_on(False)
    polygons = []
    color = []
    for ann in sorted_anns:
        m = ann['segmentation']
        img = np.ones((m.shape[0], m.shape[1], 3))
        color_mask = np.random.random((1, 3)).tolist()[0]
        for i in range(3):
            img[:,:,i] = color_mask[i]
        ax.imshow(np.dstack((img, m*0.35)))

def encode(sequence):
    count = 1
    result = []
    for x, item in enumerate(sequence): 
        if x == 0:
            continue
        elif item == sequence[x - 1]:
            count += 1
        else:        
            result.extend([count, int(sequence[x - 1])])
            count = 1            
    result.extend([count, int(sequence[len(sequence) - 1])])
    return result

def rle_encode(img_array): 
    img_array = img_array['segmentation'].astype(np.uint8)
    H, W = img_array.shape
    img_vec = ''.join([str(tmp) for tmp in img_array.flatten().tolist()])
    img_rle = encode(img_vec)
    print('rle encoding end')
    return img_rle

@app.route("/api/loraThumbnail", methods=["POST"])
def loraThumbnail():
    import requests
    json_info = request.get_json()
    lora_name = json_info['lora']
    url = f'http://localhost:{app_port_num}'  # Specify the appropriate upload endpoint URL
    for file in [lora_name + ext for ext in [".jpg", ".png", ".jpeg"]]:
        response = requests.get(url+file)
        if response.status_code == 200:
            return jsonify({"name" : "/nativefile/" + file})
    return jsonify({"name" : "https://ai-image.mogu-inc.com/20231025-180453.jpeg"})


# @cache.cached()
@app.route('/api/ping', methods=['GET'])
def ping():
    return "1"

@app.route('/api/gpuInfo', methods=['GET'])
def get_gpu_info():
    try:
        output = subprocess.check_output(['nvidia-smi', '--query-gpu=utilization.gpu,name,memory.used,memory.total', '--format=csv,nounits,noheader'])
        lines = output.decode().strip().split('\n')
        gpu_info = lines[0].split(',')

        gpu_make = gpu_info[1]  # Index 1 corresponds to GPU make
        gpu_utilization = float(gpu_info[0])  # Index 0 corresponds to GPU utilization
        vram_used = int(gpu_info[2])  # Index 2 corresponds to VRAM used
        vram_total = int(gpu_info[3])  # Index 3 corresponds to total VRAM

        gpu_info = {
            'GpuMake': gpu_make.strip(),
            'vramUsage' : f"{ round(vram_used/ 1024, 2)} / {round(vram_total/ 1024, 2)} GB",
            'GpuUtilization': f"{round(gpu_utilization, 1)} %"
        } 
        return jsonify(gpu_info)

    except (subprocess.CalledProcessError, FileNotFoundError):
        return None

@app.route('/api/imgseg-base64', methods=['POST'])
def run_all():
    start_time = time.time()
    weshop_log.info(f"start file lock on {start_time}")
    with FileLock(gpu_lock_file):
        json_info = request.get_json()
        img_url = json_info['image_url']

        image_url = img_url
        if not image_url.startswith("http://") and not image_url.startswith("https://"):
            image_url = f"http://127.0.0.1:{app_port_num}" + image_url
        weshop_log.info(f"start a subprocess doing segment on image {image_url}")

        os.chdir(os.path.dirname(__file__))
        command = [sys.executable, "run_sam.py", image_url]
        with open("segment.log", "a+") as f:
            p = subprocess.run(command, shell=True, stdout=f, stderr=f)

        with open("output/sam_output.json", "r") as file:
            json_data = json.load(file) 
        weshop_log.info(f"begin to release a file lock on {time.time()}")
        return json_data

@app.route('/api/imgto4k', methods=['GET'])
def run_codeformer():
    with FileLock(gpu_lock_file):
        image_url = request.args.get('url')
        if not image_url.startswith("http://") and not image_url.startswith("https://"):
            image_url = f"http://127.0.0.1:{app_port_num}" + image_url
        os.chdir(os.path.dirname(__file__))
        with open("realesrgan.log", "a+") as f:
            command = [sys.executable, "run_realesrgan.py", image_url]
            output = subprocess.run(command, shell=True, stdout=f, stderr=f)

        with open("realesrgan.log", "a+") as f:
            command = [sys.executable, "run_codeformer.py", image_url]
            output = subprocess.run(command, shell=True,  stdout=f, stderr=f)

        with open("output/codeformer_output.json", "r") as file:
            json_data = json.load(file) 
        return json_data


@app.route('/api/sd-models', methods=["GET"])
def list_sd_models():

    return [{
        "title": model,
        "model_name": os.path.splitext(model)[0],
        "hash": "",
        "filename": os.path.join(sd_models_path, model),
        "config": ""
        } for model in os.listdir(sd_models_path) if os.path.splitext(model)[1] in [".ckpt", ".safetensors"]]

@app.route('/api/sd-vae', methods=["GET"])
def list_sd_vae():
    return [{
        "model_name" : model,
        "filename" : os.path.join(sd_vae_path, model)
        } for model in os.listdir(sd_vae_path) if os.path.splitext(model)[1] in [".ckpt", ".safetensors"]]

@app.route('/api/sd-loras', methods=["GET"])
def list_sd_lora():
    result = []
    for model in [model for model in os.listdir(sd_lora_path) if os.path.splitext(model)[1] in [".ckpt", ".safetensors"]]:
        model_no_ext = os.path.splitext(model)[0]
        lora_dict = {
            "name" : os.path.splitext(model)[0],
            "alias": os.path.splitext(model)[0],
            "path" : os.path.join(sd_lora_path, model_no_ext),
        }
        if os.path.exists(os.path.join(sd_lora_path, model_no_ext + ".jpg")):
            lora_dict["image"] = f"/nativefile/lora_cover/{model_no_ext}.jpg"
        elif os.path.exists(os.path.join(sd_lora_path, model_no_ext + ".jpeg")):
            lora_dict["image"] = f"/nativefile/lora_cover/{model_no_ext}.jpeg"
        elif os.path.exists(os.path.join(sd_lora_path, model_no_ext + ".png")):
            lora_dict["image"] = f"/nativefile/lora_cover/{model_no_ext}.png"
        else:
            lora_dict["image"] = "https://ai-image.weshop.com/20231025-180453.jpeg"
        result.append(lora_dict)
    return result

@app.route('/api/prompt', methods=["GET"])
def prompt():
    text = request.args.get('text')
    start_time = time.time()
    candidate = []
    for key in data_dict.keys():
        if key.startswith(text):
            candidate.append(key)
    candidate.sort(key=lambda x: -data_dict[x])
    time_taken = time.time() - start_time
    print(f"time taken for {text} prompt is {time_taken}")
    return {"data" : candidate[:10]}


def construct_mask_rle(image):
    gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    _, binary_mask = cv2.threshold(gray, 220, 255, cv2.THRESH_BINARY)
    rle_mask = maskUtils.encode(np.asfortranarray(binary_mask))
    return rle_mask


def run_sam(image_url):
    from imgsam import SegmentAnything
    from urllib.request import urlopen

    response = urlopen(image_url)
    img_data = response.read()
    image = Image.open(BytesIO(img_data)).convert('RGB')
    img_array = np.array(image)
    try:
        device = torch.device('cuda')
        segment_model = SegmentAnything(device, arch="vit_h", task = 'online', output_mode='uncompressed_rle')
        anns = segment_model.generate_mask(img_array)
        anns = sorted(anns, key=lambda item: item['area'])     
        for info in anns:
            info['segmentation']['cums'] = np.cumsum(info['segmentation']['counts']).tolist()
            info['image_url'] = image_url
        json_string = json.dumps({"data" : anns, "success" : True , "message" : "success",})

        # Save JSON string to the cleared file
        with open('output/sam_output.json', 'w') as file:
            file.write(json_string)

        torch.cuda.empty_cache()

    except Exception as e:
        torch.cuda.empty_cache()
    finally:
        print(1)


def calculate_iou_mask(mask1, mask2):
    # Create binary masks
    binary_mask1 = maskUtils.decode(mask1)
    binary_mask2 = maskUtils.decode(mask2)

    # Calculate intersection and union
    intersection = binary_mask1 & binary_mask2
    union = binary_mask1 | binary_mask2
    iou = intersection.sum() / float(union.sum())
    return iou

def calculate_iou_bbox(bbox1, bbox2):
    # Unpack the bounding box coordinates
    x1, y1, w1, h1 = bbox1
    x2, y2, w2, h2 = bbox2
    # Calculate the coordinates of the intersection rectangle
    x_intersection, y_intersection = max(x1, x2), max(y1, y2)
    w_intersection = min(x1 + w1, x2 + w2) - x_intersection
    h_intersection = min(y1 + h1, y2 + h2) - y_intersection

    if w_intersection <= 0 or h_intersection <= 0:
        return 0.0

    # Calculate the areas of the bounding boxes and the intersection
    area_bbox1, area_bbox2 = w1 * h1, w2 * h2
    area_intersection = w_intersection * h_intersection
    iou = area_intersection / float(area_bbox1 + area_bbox2 - area_intersection)
    return iou

def calculate_iou(bbox, mask, image_size):
    '''
    convert bounding box to bbox mask and binary mask
    '''
    bbox_binary = np.zeros(image_size, dtype=np.uint8)
    x1, y1, x2, y2 = bbox
    height = y2 - y1
    width = x2 - x1 
    area_bbox = height * width

    bbox_binary[y1: y1 + int(1*height/1), x1 : x1 + int(1*width/1)] = 1 

    # Create binary mask from run-length encoded mask
    mask_binary = maskUtils.decode(mask)

    # Calculate intersection and union
    intersection = np.logical_and(bbox_binary, mask_binary).sum()
    union = np.logical_or(bbox_binary, mask_binary).sum()
    # iou = intersection / float(union)
    iou = intersection / float(area_bbox)
    return iou



from img2tiff import jpg_to_tiff_async, jpg_to_tiff, convert_base64_to_tiff
from tiff2psd import PSDClass
from io import BytesIO
import asyncio
import logging

def replace_port(url, new_port):
    pattern = r":\d+"  # 匹配冒号和数字组成的端口部分
    replaced_url = re.sub(pattern, ":" + new_port, url)
    return replaced_url

@app.route('/download', methods=['POST'])
def download_file():
    json_data = request.json
    generatedImageUrls = json_data.get('generatedImageUrls', list)
    originalImageUrl = json_data.get('originalImageUrl')
    mask = json_data.get('mask')
    for i in range(len(generatedImageUrls)):
        if not generatedImageUrls[i].startswith("http://") and not generatedImageUrls[i].startswith("https://"):
            generatedImageUrls[i] = f"http://127.0.0.1:{app_port_num}" + generatedImageUrls[i]
        else:
            generatedImageUrls[i] = replace_port(generatedImageUrls[i], app_port_num)
    if len(originalImageUrl) != 0:
        if not originalImageUrl.startswith("http://") and not originalImageUrl.startswith("https://"):
            originalImageUrl = f"http://127.0.0.1:{app_port_num}" + originalImageUrl
        else:
            originalImageUrl = replace_port(originalImageUrl, app_port_num)

    print(f"originalImageUrl: {originalImageUrl}")
    print(f"generatedImageUrls: {generatedImageUrls}")

    # orgTiff = jpg_to_tiff(originalImageUrl)
    # a = PSDClass(orgTiff)
    # maskTif = convert_base64_to_tiff(mask)
    # a.add_layer([maskTif], "Mask")
    # genTiffs = asyncio.run(async_call(generatedImageUrls))
    # a.add_layer(genTiffs, "Generated image")

    orgTiff = jpg_to_tiff(generatedImageUrls[0])
    a = PSDClass(orgTiff)
    if originalImageUrl and len(originalImageUrl) != 0:
        orgTiff = jpg_to_tiff(originalImageUrl)
        a.add_layer([orgTiff], "Original image")
    if mask and len(mask) != 0:
        maskTif = convert_base64_to_tiff(mask.replace('data:image/png;base64,', ''))
        a.add_layer([maskTif], "Mask")
    genTiffs = asyncio.run(async_call(generatedImageUrls))
    a.add_layer(genTiffs, "Generated image")

    # with open("output.psd", 'wb') as fd2:
    #     a.psd.write(fd2)

    # Save PSD to a BytesIO object
    psd_output = BytesIO()
    a.psd.write(psd_output)
    psd_output.seek(0)

    # Create response with attachment
    response = make_response(psd_output.getvalue())
    response.headers.set('Content-Disposition', 'attachment', filename='output.psd')
    response.headers.set('Content-Type', 'application/octet-stream')

    return response
    # pass

async def async_call(generatedImageUrls):
    tiffs = await jpg_to_tiff_async(generatedImageUrls)
    return tiffs




async def run(originalImageUrl, generatedImageUrls, mask):
    orgTiff = jpg_to_tiff(generatedImageUrls[0])
    a = PSDClass(orgTiff)
    if not originalImageUrl or len(originalImageUrl) != 0:
        orgTiff = jpg_to_tiff(originalImageUrl)
        a.add_layer([orgTiff], "Original image")
    if not mask or len(mask) != 0:
        maskTif = convert_base64_to_tiff(mask.replace('data:image/png;base64,', ''))
        a.add_layer([maskTif], "Mask")
    genTiffs = asyncio.run(async_call(generatedImageUrls))
    a.add_layer(genTiffs, "Generated image")


    # Save PSD to a BytesIO object
    psd_output = BytesIO()
    a.psd.write(psd_output)
    psd_output.seek(0)
    # pass


if __name__ == '__main__':
    #start server
    app.run(host = "127.0.0.1" , port="{}".format(60030))

# nohup python imgseg.py > output.log 2>&1 &


