import os
import io
import numpy as np
from PIL import Image, ImageOps
from google.cloud import storage
import torch

class GCPReadImageNode:
    def __init__(self):
       script_dir = os.path.dirname(os.path.abspath(__file__))
       self.output_dir = os.path.normpath(os.path.join(script_dir, "temp"))
       os.makedirs(self.output_dir, exist_ok=True)

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "bucket_name": ("STRING", {"default": "my-bucket"}),
                "bucket_path": ("STRING", {"default": "some/folder"}),
                "file_name": ("STRING", {"default": "my_image.png"}),
                "gcp_service_json": ("STRING", {"default": "/path/to/service_account.json"}),
            }
        }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("image",)
    FUNCTION = "download_from_gcp"
    OUTPUT_NODE = True
    CATEGORY = "gcp_storage"
    

    def download_from_gcp(self, bucket_name, bucket_path, file_name, gcp_service_json):
        try:
            
            script_dir = os.path.dirname(os.path.abspath(__file__))
            if not os.path.isabs(gcp_service_json):
                gcp_service_json = os.path.join(script_dir, gcp_service_json)

            if not os.path.exists(gcp_service_json):
                raise FileNotFoundError(f"Credential file not found: {gcp_service_json}")

            os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = gcp_service_json

            local_file_path = os.path.join(self.output_dir, os.path.basename(file_name))
            gcp_path = f"{bucket_path}/{os.path.basename(file_name)}".strip("/")

            storage_client = storage.Client()
            bucket = storage_client.bucket(bucket_name)
            blob = bucket.blob(gcp_path)
            print(f"[GCP READ] Downloading gs://{bucket_name}/{gcp_path} → {local_file_path}")
            blob.download_to_filename(local_file_path)
            
            i = Image.open(local_file_path)
            i = ImageOps.exif_transpose(i)
            image = i.convert("RGB")
            image_np = np.array(image).astype(np.float32) / 255.0
            
            if len(image_np.shape) == 3:
                image_np = image_np[None, ...]
            
            image = torch.from_numpy(image_np)
            image = image.contiguous().float()
            
            os.remove(local_file_path)
            
            return (image,)
        except Exception as e:
            print(str(e))
            raise e


class GCPWriteImageNode:
    def __init__(self):
        self.compress_level = 4
        self.type = "output"
        script_dir = os.path.dirname(os.path.abspath(__file__))
        self.output_dir = os.path.normpath(os.path.join(script_dir, "temp"))
        os.makedirs(self.output_dir, exist_ok=True)

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "images": ("IMAGE",),
                "bucket_name": ("STRING", {"default": "my-bucket"}),
                "bucket_path": ("STRING", {"default": "some/folder"}),
                "file_name": ("STRING", {"default": "my_image.png"}),
                "gcp_service_json": ("STRING", {"default": "/path/to/service_account.json"}),
            }
        }

    RETURN_TYPES = ("IMAGE", "STRING")
    RETURN_NAMES = ("images", "error_message")
    FUNCTION = "store_image_in_gcp"
    OUTPUT_NODE = True
    CATEGORY = "gcp_storage"

    def store_image_in_gcp(self, images, bucket_name, bucket_path, file_name, gcp_service_json):
        error_message = ""
        try:
            script_dir = os.path.dirname(os.path.abspath(__file__))
            if not os.path.isabs(gcp_service_json):
                gcp_service_json = os.path.join(script_dir, gcp_service_json)

            if not os.path.exists(gcp_service_json):
                raise FileNotFoundError(f"Credential file not found: {gcp_service_json}")

            os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = gcp_service_json

            subfolder = os.path.dirname(os.path.normpath(file_name))
            full_output_folder = os.path.join(self.output_dir, subfolder)
            full_file_path = os.path.join(full_output_folder, file_name)

            results = list()
            for (batch_number, image) in enumerate(images):
                i = 255. * image.cpu().numpy()
                img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
                metadata = None
                img.save(os.path.join(full_output_folder, file_name), pnginfo=metadata, compress_level=self.compress_level)
                results.append({
                    "filename": file_name,
                    "subfolder": subfolder,
                    "type": self.type
                })
            
            storage_client = storage.Client()
            bucket = storage_client.bucket(bucket_name)
            blob = bucket.blob(f"{bucket_path}/{file_name}")
            print(f"Uploading blob to {bucket_name}/{bucket_path}/{file_name}..")
            blob.upload_from_filename(full_file_path)
            
            os.remove(full_file_path)

            return {"ui": {"images": results}}
        except Exception as e:
            error_message = f"Upload failed: {str(e)}"
            print(error_message)
            return (None, error_message)



NODE_CLASS_MAPPINGS = {
    "GCPReadImageNode": GCPReadImageNode,
    "GCPWriteImageNode": GCPWriteImageNode,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "GCPReadImageNode": "GCP: Read Image",
    "GCPWriteImageNode": "GCP: Write Images",
}