import hashlib
import struct
from PIL import Image
from typing import Optional, Dict, Iterable
from zstandard import ZstdCompressor, ZstdDecompressor
import json
import numpy as np
import jsonschema
from datetime import datetime, UTC
import io
class Gen5DecodeError(Exception):
    pass

class Gen5CorruptHeader(Gen5DecodeError):
    def __init__(self, message: str):
        super().__init__(f"Corrupt header: {message}")

class Gen5MetadataError(Gen5DecodeError):
    def __init__(self, message: str):
        super().__init__(f"Corrupt Metadata: {message}")

class Gen5ChunkError(Gen5DecodeError):
    def __init__(self, message: str):
        super().__init__(f"Corrupt Chunk: {message}")

class Gen5LatentError(Gen5ChunkError):
    def __init__(self, message: str):
        super().__init__(f"Corrupt Latent: {message}")

class Gen5ImageError(Gen5ChunkError):
    def __init__(self, message: str):
        super().__init__(f"Corrupt Image: {message}")

class Gen5FileHandler():
    #HEADER
    HEADER_FORMAT = '<4s B B H I I I I Q'  #uses little endian
    HEADER_SIZE = struct.calcsize(HEADER_FORMAT)
    JSON_SCHEMA = """{
  "$schema": "http://json-schema.org/draft-07/schema#",
  "title": "GEN5 Metadata Schema",
  "type": "object",

  "properties": {
    "gen5_metadata": {
      "type": "object",

      "properties": {

        "file_info": {
          "type": "object",
          "properties": {
            "magic":          { "type": "string", "const": "GEN5" },
            "version_major":  { "type": "integer", "minimum": 1 },
            "version_minor":  { "type": "integer", "minimum": 0 },
            "file_size":      { "type": "integer", "minimum": 0 },
            "chunk_count":    { "type": "integer", "minimum": 0 }
          },
          "required": [
            "magic",
            "version_major",
            "version_minor",
            "file_size",
            "chunk_count"
          ]
        },

        "model_info": {
          "type": "object",
          "properties": {
            "model_name": { "type": "string" },
            "version":    { "type": "string" },
            "date":       { "type": "string" },
            "prompt":     { "type": "string" },
            "tags": {
              "type": "array",
              "items": { "type": "string" }
            },

            "generation_settings": {
              "type": "object",
              "properties": {
                "seed":         { "type": "integer", "minimum": 0 },
                "steps":        { "type": "integer", "minimum": 1 },
                "sampler":      { "type": "string" },
                "cfg_scale":    { "type": "number", "minimum": 0 },
                "scheduler":    { "type": "string" },
                "eta":          { "type": "number", "minimum": 0 },
                "guidance":     { "type": "string" },
                "precision":    { "type": "string" },
                "deterministic":{ "type": "boolean" }
              },
              "required": ["seed", "steps", "sampler"]
            },

            "hardware_info": {
              "type": "object",
              "properties": {
                "machine_name": { "type": "string" },
                "os":           { "type": "string" },
                "cpu":          { "type": "string" },
                "cpu_cores":    { "type": "integer", "minimum": 1 },

                "gpu": {
                  "type": "array",
                  "items": {
                    "type": "object",
                    "properties": {
                      "name":         { "type": "string" },
                      "memory_gb":    { "type": "number", "minimum": 0 },
                      "driver":       { "type": "string" },
                      "cuda_version": { "type": "string" }
                    },
                    "required": ["name"]
                  }
                },

                "ram_gb":      { "type": "number", "minimum": 0 },
                "framework":   { "type": "string" },
                "compute_lib": { "type": "string" }
              },
              "required": ["os"]
            }
          },

          "required": [
            "model_name",
            "version",
            "date",
            "prompt",
            "tags"
          ]
        },

        "chunks": {
          "type": "array",
          "items": {
            "type": "object",
            "properties": {
              "index":             { "type": "integer", "minimum": 0 },
              "type":              { "type": "string" },
              "flags":             { "type": "string" },
              "offset":            { "type": "integer", "minimum": 0 },
              "compressed_size":   { "type": "integer", "minimum": 0 },
              "uncompressed_size": { "type": "integer", "minimum": 0 },
              "hash":              { "type": "string" },
              "extra":             { "type": "object" },
              "compressed":        { "type": "boolean" }
            },
            "required": [
              "index",
              "type",
              "flags",
              "offset",
              "compressed_size",
              "uncompressed_size",
              "hash",
              "extra",
              "compressed"
            ]
          }
        }
      },

      "required": ["file_info", "model_info", "chunks"]
    }
  },

  "required": ["gen5_metadata"]
}
"""
    
    @classmethod
    def header_init(
        cls,
        version_major = 1,
        version_minor = 0,
        flags = 0,
        chunk_table_offset = 0,
        chunk_table_size = 0,
        chunk_count = 0,
        file_size = 0,
        reserved = 0):
        """Initialize GEN5 file header.
        Returns the packed header bytes.
        Args:
            version_major (int): Major version number.
            version_minor (int): Minor version number.
            flags (int): Header flags.
            chunk_table_offset (int): Offset to the chunk table.
            chunk_table_size (int): Size of the chunk table.
            chunk_count (int): Number of chunks.
            file_size (int): Total file size.
            reserved (int): Reserved for future use.
        """
        # header init
        return struct.pack(
            cls.HEADER_FORMAT,
            b'GEN5',
            version_major,
            version_minor,
            flags,
            chunk_table_offset,
            chunk_table_size,
            chunk_count,
            file_size,
            reserved
        )
    @classmethod
    def header_parse(cls, data: bytes):
        """Parse GEN5 file header.
        Args:
            data (bytes): The header bytes to parse.
            Returns:
            dict: Parsed header fields.
            
        Returns:
            dict: Parsed header fields.
            """
        unpacked = struct.unpack(cls.HEADER_FORMAT, data)
        return {
            'magic': unpacked[0],
            'version_major': unpacked[1],
            'version_minor': unpacked[2],
            'flags': unpacked[3],
            'chunk_table_offset': unpacked[4],
            'chunk_table_size': unpacked[5],
            'chunk_count': unpacked[6],
            'file_size': unpacked[7],
            'reserved': unpacked[8]
        }

    def header_validate(self, header: dict) -> bool:
        """Validate GEN5 file header.
        Args:
            header (dict): Parsed header fields.
            Returns:
            bool: True if it is valid or False otherwise."""
        if header['magic'] != b'GEN5':
            return False
        if header['version_major'] < 1:
            return False
        if header['chunk_count'] < 0:
            return False
        return True
    @staticmethod
    def png_to_bytes(png_path: str) -> bytes:
        """Convert PNG image to bytes, preserving transparency.
        Args:
            png_path (str): Path to the PNG image.
        Returns:
            bytes: PNG image data in bytes.

        """
        try:
            with Image.open(png_path) as img:
                img = img.convert("RGBA") 
                buf = io.BytesIO()
                img.save(buf, format="PNG")
                return buf.getvalue()
        except Exception as e:
            raise Gen5ImageError(f"Failed to convert PNG to bytes: {e}") from e    @staticmethod
    def bytes_to_png(img_bytes: bytes) -> Image.Image:
        """Convert bytes back to PNG image."""
        try:
            buffer = io.BytesIO(img_bytes)
            img = Image.open(buffer)
            return img
        except Exception as e:
            raise Gen5ImageError(f"Failed to convert bytes to PNG: {e}") from e

    def _chunk_packer_binary(self, chunk_type:bytes, chunk_flags: bytes, data:bytes) -> bytes:
        """
        Pack a binary chunk with header and compress it using Zstandard.
        Args:
            chunk_type (bytes): 4-byte chunk type identifier.
            chunk_flags (bytes): 4-byte chunk flags.
            data (bytes): Payload data.
        Returns:
            bytes: Compressed chunk data.
        """
        chunk_size = len(data)
        #chunk maker
        chunk_header = struct.pack('<4s 4s I', chunk_type, chunk_flags, chunk_size)
        #compress chunk
        chunk = chunk_header + data
        compressed_chunk = ZstdCompressor().compress(chunk)
        return compressed_chunk
        
    #METADATA
    def metadata_validator(self, manifest) -> bool:
        """Validate metadata manifest using JSON Schema.
        Args:
            manifest (dict): Metadata manifest.
        Returns:
            bool: True if it is valid or False otherwise.
        Raises:
            Gen5MetadataError: If the metadata is invalid.
        """
        json_schema = self.JSON_SCHEMA
        schema = json.loads(json_schema)
        try:
            jsonschema.validate(instance=manifest, schema=schema)
            return True
        except Exception as e:
            raise Gen5MetadataError(f"Invalid metadata: {e}")
            
                
    def metadata_compressor(self, manifest):
        json_bytes = json.dumps(manifest, indent=2).encode("utf-8")
        chunk_type = b"META"
        chunk_flags = b"0000"
        chunk_size = len(json_bytes)
        header = struct.pack("<4s 4s I", chunk_type, chunk_flags, chunk_size)
        compressed = ZstdCompressor().compress(header + json_bytes)
        return compressed
    
    def build_manifest(
        self,
        version_major: int,
        version_minor: int,
        model_name: str,
        model_version: str,
        prompt: str,
        tags: list,
        chunk_records: list,
        generation_settings: Optional[dict] = None,
        hardware_info: Optional[dict] = None
    ):

        """
        chunk_records must be a list of dicts, with each having:
            {
                "type": str,
                "flags": str,
                "offset": int,
                "compressed_size": int,
                "uncompressed_size": int,
                "hash": str,
                "extra": dict
            }
        Args:
            version_major (int): Major version number.
            version_minor (int): Minor version number.
            model_name (str): Name of the model.
            model_version (str): Version of the model.
            prompt (str): Prompt used for generation.
            tags (list): List of tags.
            chunk_records (list): List of chunk record dictionaries.
        Returns:
            dict: Manifest dictionary.
        """

        manifest = {
            "gen5_metadata": {
                "file_info": {
                    "magic": "GEN5",
                    "version_major": version_major,
                    "version_minor": version_minor,
                    "file_size": 0,
                    "chunk_count": len(chunk_records)
                },

                "model_info": {
                    "model_name": model_name,
                    "version": model_version,
                    "date": datetime.now(UTC).isoformat(),
                    "prompt": prompt,
                    "tags": tags,

                    "generation_settings": generation_settings or {
                        "seed": 0,
                        "steps": 0,
                        "sampler": "",
                        "cfg_scale": 0.0,
                        "scheduler": "",
                        "eta": 0.0,
                        "guidance": "",
                        "precision": "fp16",
                        "deterministic": True
                    },

                    "hardware_info": hardware_info or {
                        "machine_name": "",
                        "os": "",
                        "cpu": "",
                        "cpu_cores": 0,
                        "gpu": [],
                        "ram_gb": 0.0,
                        "framework": "",
                        "compute_lib": ""
                    }
                },

                
                
                "chunks": []
                
            }
        }

        #build chunk list with indexes
        for idx, rec in enumerate(chunk_records):
            manifest["gen5_metadata"]["chunks"].append({
                "index": idx,
                "type": rec["type"],
                "flags": rec["flags"],
                "offset": rec["offset"],
                "compressed_size": rec["compressed_size"],
                "uncompressed_size": rec["uncompressed_size"],
                "hash": rec["hash"],
                "extra": rec.get("extra", {}),
                "compressed": rec.get("compressed", True)
            })

        
        return manifest
        

    def metadata_parser(self, compressed_chunk: bytes) -> dict:
        """Parse and decompress metadata chunk.
        Args:
            compressed_chunk (bytes): Compressed metadata chunk.
        Returns:
            dict: Parsed metadata manifest."""
        decompressor = ZstdDecompressor()
        chunk = decompressor.decompress(compressed_chunk)
        chunk_type, chunk_flags, chunk_size = struct.unpack('<4s 4s I', chunk[:12])
        json_bytes = chunk[12:12+chunk_size]
        manifest = json.loads(json_bytes.decode("utf-8"))
        return manifest

    #IMAGE-DATA
    def image_bytes_validator(self, image_bytes: bytes):
        """Validate image bytes.
        Args:
            image_bytes (bytes): Image bytes to validate.
        """
        try:
            img = io.BytesIO(image_bytes)
            with Image.open(img) as img:
                img.verify()
            return True
        except Exception as e:
            raise Gen5ImageError(f"Invalid image bytes: {e}")


    def image_data_chunk_builder(self, image_binary: bytes):
        self.image_bytes_validator(image_binary)
        try:
            chunk_type = b'DATA'
            chunk_flags = b'0000'
            chunk_size = len(image_binary)
            chunk_header = struct.pack('<4s 4s I', chunk_type, chunk_flags, chunk_size)
            chunk = chunk_header + image_binary
            compressed_chunk = ZstdCompressor().compress(chunk)
            return compressed_chunk         
        except Exception as e:
            raise Gen5ImageError(f"Failed to build image data chunk: {e}") from e

    def image_data_chunk_parser(self, compressed_chunk):
        """Parse image data chunk.
        Args:
            compressed_chunk (bytes): Compressed image data chunk.
        Returns:
            dict: Parsed image data chunk information.
        """
        try:
            decompressor = ZstdDecompressor()
            chunk = decompressor.decompress(compressed_chunk)
            chunk_type, chunk_flags, chunk_size = struct.unpack('<4s 4s I', chunk[:12])
            image_data = chunk[12:12+chunk_size]
            return {
                "chunk_type": chunk_type,
                "chunk_flags": chunk_flags,
                "chunk_size": chunk_size,
                "image_data": image_data
            }
        except Exception as e:
            raise Gen5ImageError(f"Failed to parse image data chunk: {e}") from e

    #LATENT CHUNK
    def latent_packer(self, latent: Dict[str, np.ndarray], file_offset: int = 0, chunk_records=None, should_compress: bool = True, convert_float16: bool = True) -> list:
        """Pack latent arrays into compressed chunk.
        Args:
            latent (Dict[str, np.ndarray]): Dictionary of latent arrays.
            file_offset (int): Current file offset for chunk placement.
            chunk_records (list): List to append chunk records to.
        Returns:
            list: List of compressed latent chunks."""
        latents=[]
        if chunk_records is None:
            chunk_records = []
        for latent_array in latent.values():
            if latent_array.dtype not in (np.float16, np.float32):
                raise ValueError("Latent must be float16 or float32")
            
            if convert_float16:
                if latent_array.dtype != np.float16:
                    latent_array = latent_array.astype(np.float16)
                chunk_flags = b"F16"
            else:
                if latent_array.dtype == np.float32:
                    chunk_flags = b'F32'
                elif latent_array.dtype == np.float16:
                    chunk_flags = b'F16'
            chunk_type = b'LATN'
            data_bytes = latent_array.tobytes()
            uncompressed_size = len(data_bytes)
            if should_compress == True:
                #compress the latent
                compressed = self._chunk_packer_binary(chunk_type, chunk_flags, data_bytes)
                compressed_size = len(compressed)
            else:
                compressed = data_bytes
                compressed_size = uncompressed_size

            #create manifest
            chunk_records.append({
                "type": "LATN",
                "flags": chunk_flags.decode('ascii'),
                "offset": file_offset,
                "compressed_size": compressed_size,
                "compressed": should_compress,
                "uncompressed_size": uncompressed_size,
                "hash": hashlib.sha256(data_bytes).hexdigest(),
                "extra": {
                    "shape": list(latent_array.shape),
                    "dtype": str(latent_array.dtype)
                }
            })
            file_offset += compressed_size
            latents.append(compressed)
            
        return latents

    def latent_parser(self, compressed_chunk: bytes, shape: tuple, compressed: bool = True):
        """Parse a given single latent chunk and return the latent array.
        
        Args:
            compressed_chunk (bytes): Compressed latent chunk.
            shape (tuple): Shape to reshape the latent array to.
            
        Returns:
            np.ndarray: Parsed latent array.
        """
        if compressed == True:
            decompressor = ZstdDecompressor()
            try:
                decompressed = decompressor.decompress(compressed_chunk)
            except Exception as e:
                raise Gen5LatentError(f"Failed to decompress latent chunk: {e}")
            if len(decompressed) < 12:
                raise Gen5LatentError("Truncated latent chunk")

            chunk_type, chunk_flags, chunk_size = struct.unpack('<4s 4s I', decompressed[:12])
            data_bytes = decompressed[12:12 + chunk_size]

            if len(data_bytes) != chunk_size:
                raise Gen5LatentError("Truncated latent chunk")

            flag_str = chunk_flags.decode("utf-8", errors="ignore").replace('\x00', '').strip()
            if flag_str == "F16":
                dtype = np.float16
            elif flag_str == "F32":
                dtype = np.float32
            else:
                raise ValueError(f"Unknown latent dtype flag: {flag_str!r}")

            print("chunk_size:", chunk_size)
            print("len(data_bytes):", len(data_bytes))
            print("expected bytes:", np.prod(shape) * np.dtype(dtype).itemsize)
            expected = np.prod(shape) * np.dtype(dtype).itemsize
            if len(data_bytes) != expected:
                raise Gen5LatentError(
                    f"Size mismatch: expected {expected} bytes, got {len(data_bytes)}"
                    )
            arr = np.frombuffer(data_bytes, dtype=dtype).copy()
            arr = arr.reshape(shape)
            return arr
        else:
            if len(compressed_chunk) < 12:
                raise Gen5LatentError("Truncated latent chunk")

            chunk_type, chunk_flags, chunk_size = struct.unpack(
                "<4s 4s I", compressed_chunk[:12]
            )

            data_bytes = compressed_chunk[12:12 + chunk_size]

            if len(data_bytes) != chunk_size:
                raise Gen5LatentError("Truncated latent chunk")

            flag_str = chunk_flags.decode("utf-8", errors="ignore").strip()
            dtype = np.float16 if flag_str == "F16" else np.float32

            arr = np.frombuffer(data_bytes, dtype=dtype).copy()
            return arr.reshape(shape)

    #RUNNING

    def file_encoder(self, filename: str, latent: Dict[str, np.ndarray], chunk_records: list,
                    model_name: str, model_version: str, prompt: str, tags: list, img_binary: bytes, should_compress: bool = True, convert_float16: bool = True, generation_settings: Optional[dict] = None, hardware_info: Optional[dict] = None):
        """ Orchestrator function to encode GEN5 file.
        Args:
            filename (str): Output GEN5 filename. [.gen5 extension is REQUIRED]
            latent (Dict[str, np.ndarray]): Dictionary of latent arrays.
            chunk_records (list): List to append chunk records to.
            model_name (str): Name of the model.
            model_version (str): Version of the model.
            prompt (str): Prompt used for generation.
            tags (list): List of tags.
            img_binary (bytes): Binary image data.
            Returns:
            dict: Dictionary with header, latent chunks, metadata, and image chunk.
            """
        # Pack the latent chunks first and update chunk_records with correct offsets
        if not filename.endswith('.gen5'):
            raise ValueError("Filename must have a .gen5 extension")

        current_offset = self.HEADER_SIZE
        latent_chunks = self.latent_packer(latent, file_offset=current_offset, chunk_records=chunk_records, should_compress=should_compress, convert_float16=convert_float16)
        latent_chunk = b"".join(latent_chunks)
        current_offset += len(latent_chunk)

        #Pack image chunk if provided
        image_chunk = None
        if img_binary is not None:
            image_chunk = self.image_data_chunk_builder(img_binary)
            image_chunk_record = {
                "type": "DATA",
                "flags": "0000",
                "offset": current_offset,
                "compressed_size": len(image_chunk),
                "uncompressed_size": len(img_binary),
                "hash": hashlib.sha256(img_binary).hexdigest(),
                "extra": {},
                "compressed": True
            }
            chunk_records.append(image_chunk_record)
            current_offset += len(image_chunk)
            print("ENCODER STORED FLAG:", chunk_records[-1]["flags"])


        #build manifest with all chunks
        manifest = self.build_manifest(
            version_major=1,
            version_minor=0,
            model_name=model_name,
            model_version=model_version,
            prompt=prompt,
            tags=tags,
            chunk_records=chunk_records,
            generation_settings=generation_settings,
            hardware_info=hardware_info
        )
        #ensure that manifest/metadata is valid
        self.metadata_validator(manifest)
        #compress metadata
        compressed_metadata = self.metadata_compressor(manifest)
        metadata_size = len(compressed_metadata)
        
        #calculate total file size
        total_file_size = self.HEADER_SIZE + len(latent_chunk) + (len(image_chunk) if image_chunk else 0) + metadata_size

        #update file_size in manifest
        manifest["gen5_metadata"]["file_info"]["file_size"] = total_file_size
        
        #recompress metadata with the updated file_size
        compressed_metadata = self.metadata_compressor(manifest)
        metadata_size = len(compressed_metadata)
        total_file_size = self.HEADER_SIZE + len(latent_chunk) + (len(image_chunk) if image_chunk else 0) + len(compressed_metadata)

        #final header
        header = self.header_init(
            version_major=1,
            version_minor=0,
            flags=0,
            chunk_table_offset=self.HEADER_SIZE + len(latent_chunk) + (len(image_chunk) if image_chunk else 0),
            chunk_table_size=metadata_size,
            chunk_count=len(chunk_records),
            file_size=total_file_size
        )

        
        with open(filename, "wb") as f:
            f.write(header)
            f.write(latent_chunk)
            if image_chunk is not None:
                f.write(image_chunk)
            f.write(compressed_metadata)

        return {
            "header": header,
            "latent_chunks": latent_chunk,
            "metadata_chunk": compressed_metadata,
            "image_chunk": image_chunk
        }


    def file_decoder(self, filename: str):
        """ Orchestrator function to decode GEN5 file.
        Args:
            filename (str): Input GEN5 filename.
        Returns:
            dict: Dictionary with header, chunks, and metadata.
        """
        with open(filename, "rb") as f:
            header_bytes = f.read(self.HEADER_SIZE)
            header = self.header_parse(header_bytes)
            if not self.header_validate(header):
                raise Gen5CorruptHeader(message=f"Invalid header: {header}")
            f.seek(header['chunk_table_offset'])
            metadata_compressed = f.read(header['chunk_table_size'])
            metadata = self.metadata_parser(metadata_compressed)
            chunk_records = metadata['gen5_metadata']['chunks']
            chunks = {}
            chunks['latent'] = []
            
            for record in chunk_records:
                chunk_type = record['type']
                compressed = record.get("compressed", True)

                f.seek(record['offset'])
                raw_chunk = f.read(record['compressed_size'])

                if len(raw_chunk) != record['compressed_size']:
                    raise Gen5ChunkError(f"Truncated chunk {chunk_type} at offset {record['offset']}")

                if chunk_type == "LATN":
                    shape = tuple(record['extra']['shape'])

                    if compressed:
                        print("decoded raw_chunk len:", len(raw_chunk))
                        print("expected total bytes:", np.prod(shape))
                        latent_array = self.latent_parser(raw_chunk, shape, True)
                    else:
                        latent_array = self.latent_parser(raw_chunk, shape, False)

                    chunks['latent'].append(latent_array)
                elif chunk_type == "DATA":
                    if compressed:
                        parsed = self.image_data_chunk_parser(raw_chunk)
                        chunks['image'] = parsed['image_data']
                    else:
                        # DATA chunks have the same header layout even when not compressed
                        chunk_type_b, flags_b, size = struct.unpack('<4s 4s I', raw_chunk[:12])
                        chunks['image'] = raw_chunk[12:12+size]
                else:
                    raise ValueError(f"Unknown chunk type: {chunk_type}. Supported: LATN, DATA")

            
            return {
                "header": header,
                "chunks": chunks,
                "metadata": metadata
            }

    def make_lazy_latent_loader(self, filename: str, chunk_record: dict):
        """Return a callable that loads the latent array on demand."""
        loaded_array = None

        def load():
            nonlocal loaded_array
            if loaded_array is None:
                with open(filename, "rb") as f:
                    try:
                        offset = chunk_record['offset']
                        size = chunk_record['compressed_size']
                        
                        shape = tuple(chunk_record['extra']['shape'])
                        compressed = chunk_record.get("compressed", True)
                        f.seek(0, 2)                 # move to end
                        file_size = f.tell()
                        if offset + size > file_size:
                            raise Gen5LatentError(
                                f"Chunk at offset {offset} with size {size} exceeds file bounds ({file_size})."
                            )
                        f.seek(offset)
                        compressed_chunk = f.read(size)
                        loaded_array = self.latent_parser(compressed_chunk, shape, compressed)
                    except Exception as e:
                        raise Gen5LatentError(f"Failed to load the latent chunk : {chunk_record} | {e}")

            return loaded_array
        return load

    def iter_lazy_latents(self, filename: str, chunk_records: list):
        """Yield callables for each latent chunk for lazy loading."""
        for record in chunk_records:
            if record['type'] == "LATN":
                yield self.make_lazy_latent_loader(filename, record)
