import SimpleITK as sitk
import os
from datetime import datetime
from core.cylinder import generate_cylinder_numpy

def save_cyl(best_position_l, best_position_r, spacing, IMAGE3, output_base, CBT=True):
    
    image3 = sitk.ReadImage(IMAGE3) #Reading image2 or image3 can get the same shape and origin actually 
    image3_array = sitk.GetArrayFromImage(image3)
    image3_shape = image3_array.shape
    origin = image3.GetOrigin()
    diameter_l = best_position_l[5]
    diameter_r = best_position_r[5]
    length_l = best_position_l[6]
    length_r = best_position_r[6]

    cyl_l = generate_cylinder_numpy(diameter_l, 
                                    length_l, 
                                    best_position_l[0], 
                                    best_position_l[1], 
                                    best_position_l[2], 
                                    best_position_l[3], 
                                    best_position_l[4], 
                                    image3_shape, 
                                    spacing)
    
    cyl_r = generate_cylinder_numpy(diameter_r, 
                                    length_r, 
                                    best_position_r[0], 
                                    best_position_r[1], 
                                    best_position_r[2], 
                                    best_position_r[3], 
                                    best_position_r[4], 
                                    image3_shape, 
                                    spacing)

    cyl_L = sitk.GetImageFromArray(cyl_l)
    cyl_R = sitk.GetImageFromArray(cyl_r)

    cyl_L.SetOrigin(origin)
    cyl_R.SetOrigin(origin)

    image_path = IMAGE3
    date_str = datetime.now().strftime("%Y%m%d")
    patient_id = os.path.basename(os.path.dirname(image_path))
    file_name = os.path.basename(image_path) 
    level = file_name.split('_')[0]
    output_folder = os.path.join(output_base, date_str, patient_id)

    os.makedirs(output_folder, exist_ok=True)

    def get_unique_filename(path):

        if not os.path.exists(path):
            return path

        # 特別處理 .nii.gz 檔案
        if path.endswith(".nii.gz"):
            base = path[:-7]  # 去除 .nii.gz
            ext = ".nii.gz"
        else:
            base, ext = os.path.splitext(path)

        i = 1
        new_path = f"{base}_{i}{ext}"
        while os.path.exists(new_path):
            i += 1
            new_path = f"{base}_{i}{ext}"
        return new_path

    if CBT:
        output_path_l = get_unique_filename(f'{output_folder}/{level}_{diameter_l}_{length_l}_CBT_L.nii.gz')
        output_path_r = get_unique_filename(f'{output_folder}/{level}_{diameter_r}_{length_r}_CBT_R.nii.gz')
        sitk.WriteImage(cyl_L, output_path_l)
        sitk.WriteImage(cyl_R, output_path_r)

    else:
        output_path_l = get_unique_filename(f'{output_folder}/{level}_{diameter_l}_{length_l}_TPS_L.nii.gz')
        output_path_r = get_unique_filename(f'{output_folder}/{level}_{diameter_r}_{length_r}_TPS_R.nii.gz')
        sitk.WriteImage(cyl_L, output_path_l)
        sitk.WriteImage(cyl_R, output_path_r)

    return output_path_l, output_path_r