#!/usr/bin/env python3
"""
Minimal NumPy-based image utilities required by tp2.py.

Supported formats: PGM/PPM (ASCII P2/P3 and RAW P5/P6).
"""

import numpy as np
from typing import Dict

# ------------------------------
# Constants and format mapping
# ------------------------------
PGM_RAW = 0
PGM_ASCII = 1
PPM_RAW = 2
PPM_ASCII = 3

FILE_MAGIC_TO_FORMAT = {
    'P5': PGM_RAW,
    'P2': PGM_ASCII,
    'P6': PPM_RAW,
    'P3': PPM_ASCII,
}


# ------------------------------
# Core helpers
# ------------------------------

def fatal_error(message: str) -> None:
    print(message)
    raise SystemExit(1)


def allocate_image(width: int,
                   height: int,
                   channels: int,
                   dtype=np.float32) -> Dict:
    if width <= 0 or height <= 0 or channels not in (1, 3):
        fatal_error('allocate_image: invalid')
    data = np.zeros((height, width, channels), dtype=dtype)
    return {
        'width': width,
        'height': height,
        'channels': channels,
        'data': data,
    }


def same_size(a: Dict, b: Dict) -> bool:
    return (a['width'] == b['width'] and
            a['height'] == b['height'] and
            a['channels'] == b['channels'])


def clone_image(img: Dict) -> Dict:
    out = allocate_image(img['width'],
                         img['height'],
                         img['channels'],
                         dtype=img['data'].dtype)
    out['data'][:] = img['data']
    return out


# ------------------------------
# I/O: Load and Save PNM (PGM/PPM)
# ------------------------------

def _read_pnm_header(f):
    """Read magic + width + height + maxval, skipping '#' comments."""
    magic = f.readline().strip()
    magic = magic.decode('ascii') if isinstance(magic, bytes) else magic
    if magic not in FILE_MAGIC_TO_FORMAT:
        fatal_error('LoadImage Error: Unsupported file type')
    fmt = FILE_MAGIC_TO_FORMAT[magic]

    def next_token():
        token = []
        while True:
            ch = f.read(1)
            if not ch:
                break
            if ch == b'#' or ch == '#':
                f.readline()
                continue
            is_space = (ch.isspace() if isinstance(ch, str)
                        else ch.isspace())
            if is_space:
                if token:
                    break
                continue
            token.append(ch)
        if not token:
            return None
        bts = b''.join(c if isinstance(c, bytes)
                       else c.encode('ascii') for c in token)
        return bts.decode('ascii')

    width = int(next_token())
    height = int(next_token())
    maxval = int(next_token())
    return magic, fmt, width, height, maxval


def LoadImage(file_name: str) -> Dict:
    if not file_name:
        fatal_error('Error: Please specify a filename')

    with open(file_name, 'rb') as f:
        magic, fmt, width, height, maxval = _read_pnm_header(f)
        channels = 1 if magic in ('P2', 'P5') else 3
        img = allocate_image(width, height, channels)

        if fmt == PGM_ASCII:
            # reopen as text to stream integers
            with open(file_name, 'r', encoding='ascii') as fa:
                # skip header lines and comments
                line = fa.readline()
                tokens = []
                while len(tokens) < 3:
                    line = fa.readline()
                    if not line:
                        break
                    if line.strip().startswith('#'):
                        continue
                    tokens.extend(line.strip().split())
                values = []
                for line in fa:
                    if line.strip().startswith('#'):
                        continue
                    values.extend(line.strip().split())
                arr = np.array(values, dtype=np.float32)
                need = width * height
                if arr.size < need:
                    fatal_error('PGM_ASCII: unexpected end of file')
                arr = arr[:need].reshape((height, width))
                img['data'][:, :, 0] = arr * (255.0 / maxval)
            return img

        if fmt == PGM_RAW:
            raw = np.frombuffer(f.read(width * height), dtype=np.uint8)
            if raw.size < width * height:
                fatal_error('PGM_RAW: unexpected end of file')
            img['data'][:, :, 0] = raw.reshape((height, width)) * (
                255.0 / maxval
            )
            return img

        if fmt == PPM_ASCII:
            with open(file_name, 'r', encoding='ascii') as fa:
                line = fa.readline()
                tokens = []
                while len(tokens) < 3:
                    line = fa.readline()
                    if not line:
                        break
                    if line.strip().startswith('#'):
                        continue
                    tokens.extend(line.strip().split())
                values = []
                for line in fa:
                    if line.strip().startswith('#'):
                        continue
                    values.extend(line.strip().split())
                arr = np.array(values, dtype=np.float32)
                need = width * height * 3
                if arr.size < need:
                    fatal_error('PPM_ASCII: unexpected end of file')
                arr = arr[:need].reshape((height, width, 3))
                img['data'][:, :, :] = arr * (255.0 / maxval)
            return img

        if fmt == PPM_RAW:
            raw = np.frombuffer(f.read(width * height * 3), dtype=np.uint8)
            if raw.size < width * height * 3:
                fatal_error('PPM_RAW: unexpected end of file')
            img['data'][:, :, :] = raw.reshape((height, width, 3)) * (
                255.0 / maxval
            )
            return img

        fatal_error('Unsupported format')


def SaveImage(img: Dict, file_name: str, ff: int) -> None:
    if not file_name:
        fatal_error('Error!! SaveImage')

    H, W, C = img['data'].shape
    data_u8 = np.clip(img['data'], 0, 255).astype(np.uint8)

    if ff == PGM_ASCII:
        with open(file_name, 'w', encoding='ascii') as out:
            out.write(f"P2\n{W} {H}\n255\n")
            for y in range(H):
                row = ' '.join(str(int(v)) for v in data_u8[y, :, 0])
                out.write(row + "\n")
        return

    if ff == PGM_RAW:
        with open(file_name, 'wb') as out:
            out.write(f"P5\n{W} {H}\n255\n".encode('ascii'))
            out.write(data_u8[:, :, 0].tobytes())
        return

    if ff == PPM_ASCII:
        with open(file_name, 'w', encoding='ascii') as out:
            out.write(f"P3\n{W} {H}\n255\n")
            arr = (data_u8 if C > 1 else
                   np.repeat(data_u8[:, :, 0:1], 3, axis=2))
            for y in range(H):
                row_vals = arr[y].reshape(-1, 3)
                row = ' '.join(
                    f"{int(r)} {int(g)} {int(b)}"
                    for r, g, b in row_vals
                )
                out.write(row + "\n")
        return

    if ff == PPM_RAW:
        with open(file_name, 'wb') as out:
            out.write(f"P6\n{W} {H}\n255\n".encode('ascii'))
            arr = (data_u8 if C > 1 else
                   np.repeat(data_u8[:, :, 0:1], 3, axis=2))
            out.write(arr.tobytes())
        return

    fatal_error('save_image: unsupported FILE_FORMAT')

