from __future__ import annotations

from pathlib import Path
import json
import laspy
import numpy as np


def structurize_las(las_path: Path, output_dir: Path, progress_callback=None) -> Path:
    """원본 LAS/LAZ 파일을 읽어 전체 속성을 포함한 고속 structured.npz 파일로 구조화합니다.

    보존 필드:
      - xyz (float64): 실제 좌표값 (scale, offset 적용됨)
      - intensity (uint16): 반사 강도
      - rgb (uint8): RGB 색상 (8bit 정규화 처리)
      - return_number (uint8): 반사 순서
      - number_of_returns (uint8): 총 반사 횟수
      - classification (uint8): 원본 클래스 (없는 경우 기본값 0)
    """
    output_dir.mkdir(parents=True, exist_ok=True)
    npz_output_path = output_dir / "structured.npz"

    # 1. 헤더 정보 및 속성 가용성 조사
    with laspy.open(las_path) as las_file:
        header = las_file.header
        total_points = int(header.point_count)
        x_min, y_min, z_min = float(header.mins[0]), float(header.mins[1]), float(header.mins[2])
        x_max, y_max, z_max = float(header.maxs[0]), float(header.maxs[1]), float(header.maxs[2])
        
        point_format = header.point_format
        has_rgb = all(name in point_format.dimension_names for name in ("red", "green", "blue"))
        has_intensity = "intensity" in point_format.dimension_names
        has_returns = "return_number" in point_format.dimension_names and "number_of_returns" in point_format.dimension_names

    # 2. 데이터 청크 로드 및 구조화
    xyz_list = []
    intensity_list = []
    rgb_list = []
    ret_num_list = []
    num_rets_list = []
    cls_list = []
    
    # 50만 포인트 단위 청크 이터레이션으로 대용량 대응
    loaded_points = 0
    with laspy.open(las_path) as las_file:
        for chunk in las_file.chunk_iterator(500_000):
            # XYZ 좌표 (정밀한 연산을 위해 float64 권장)
            xs = np.asarray(chunk.x, dtype=np.float64)
            ys = np.asarray(chunk.y, dtype=np.float64)
            zs = np.asarray(chunk.z, dtype=np.float64)
            xyz_list.append(np.stack([xs, ys, zs], axis=1))

            # 강도(Intensity)
            if has_intensity:
                intensity_list.append(np.asarray(chunk.intensity, dtype=np.uint16))
            else:
                intensity_list.append(np.zeros(len(xs), dtype=np.uint16))

            # RGB
            if has_rgb:
                r = np.asarray(chunk.red, dtype=np.float32)
                g = np.asarray(chunk.green, dtype=np.float32)
                b = np.asarray(chunk.blue, dtype=np.float32)
                
                # 16-bit RGB 대응 (최댓값이 255보다 큰 경우 8-bit로 스케일링)
                max_val = max(r.max() if r.size else 0, g.max() if g.size else 0, b.max() if b.size else 0)
                if max_val > 255.0:
                    r = r / 256.0
                    g = g / 256.0
                    b = b / 256.0
                
                rgb_stacked = np.stack([r, g, b], axis=1).clip(0, 255).astype(np.uint8)
                rgb_list.append(rgb_stacked)
            else:
                # RGB가 없는 경우 디폴트 Gray값(128)으로 채움
                rgb_list.append(np.full((len(xs), 3), 128, dtype=np.uint8))

            # Return 정보 (반사 횟수)
            if has_returns:
                ret_num_list.append(np.asarray(chunk.return_number, dtype=np.uint8))
                num_rets_list.append(np.asarray(chunk.number_of_returns, dtype=np.uint8))
            else:
                ret_num_list.append(np.ones(len(xs), dtype=np.uint8))
                num_rets_list.append(np.ones(len(xs), dtype=np.uint8))

            # 분류 정보(Classification)
            if "classification" in point_format.dimension_names:
                cls_list.append(np.asarray(chunk.classification, dtype=np.uint8))
            else:
                cls_list.append(np.zeros(len(xs), dtype=np.uint8))

            loaded_points += len(xs)
            if progress_callback:
                pct = int((loaded_points / total_points) * 100) if total_points else 100
                progress_callback(pct)



    # 3. 배열 합치기 및 압축 직렬화
    xyz = np.concatenate(xyz_list, axis=0)
    intensity = np.concatenate(intensity_list, axis=0)
    rgb = np.concatenate(rgb_list, axis=0)
    ret_num = np.concatenate(ret_num_list, axis=0)
    num_rets = np.concatenate(num_rets_list, axis=0)
    classification = np.concatenate(cls_list, axis=0)

    np.savez_compressed(
        npz_output_path,
        xyz=xyz,
        intensity=intensity,
        rgb=rgb,
        return_number=ret_num,
        number_of_returns=num_rets,
        classification=classification,
        bounds=np.array([[x_min, x_max], [y_min, y_max], [z_min, z_max]], dtype=np.float64),
        total_points=np.array([total_points], dtype=np.int64),
        has_rgb=np.array([1 if has_rgb else 0], dtype=np.int8)
    )

    return npz_output_path