import os, mmcv, cv2, json
import numpy as np
from pathlib import Path
from PIL import Image
from tqdm import tqdm
from mmseg.apis import init_segmentor, inference_segmentor
from shapely.geometry import Polygon, mapping
import sys

_DETECT_DIR = Path(__file__).parent      # mx15hdi/Detect/
_MX15HDI_DIR = _DETECT_DIR.parent       # mx15hdi/


def load_model():
    """서버 시작 시 1회 호출. 로드된 모델 객체를 반환한다."""
    config = str(_DETECT_DIR / 'V7_SPECIAL.py')
    checkpoint = str(_DETECT_DIR / 'epoch_165.pth')
    model = init_segmentor(config, checkpoint, device='cuda:0')
    model.PALETTE = [
        [0, 0, 0],       # background
        [0, 0, 204],     # black
        [180, 180, 180], # brown
        [255, 255, 0],   # rainbow
        [178, 102, 255]  # silver
    ]
    return model


def blend_images(original_img, color_mask, alpha=0.6):
    """
    Blend original image and color mask with alpha transparency.
    Inputs: numpy arrays HWC uint8
    """
    blended = cv2.addWeighted(original_img, 1 - alpha, color_mask, alpha, 0)
    return blended


def run_inference(model, file_id: str, write_files: bool = False) -> dict:
    """
    사전 로드된 모델로 file_id 폴더 내 이미지를 세그멘테이션한다.

    Args:
        model: load_model()로 로드된 모델 객체.
        file_id: 처리할 이미지 폴더명.
        write_files: True이면 Detect/result/ 와 Detect/Mask_result/ 에 중간 파일 저장.
                     False이면 디스크 쓰기 생략 (기본값).

    Returns:
        inference_cache: {image_filename: {'blended': ndarray, 'mask': ndarray, 'ext': str}}
        이 값을 run_georeference()에 전달하면 중간 파일 읽기를 생략할 수 있다.
    """
    img_path = str(_MX15HDI_DIR / 'Metadata' / 'Image' / 'Original_Images' / file_id)
    output_folder = str(_DETECT_DIR / 'result' / file_id)
    mask_folder = str(_DETECT_DIR / 'Mask_result' / file_id)

    if not os.path.exists(img_path):
        raise FileNotFoundError(f"이미지 폴더가 존재하지 않습니다: {img_path}")

    if write_files:
        os.makedirs(output_folder, exist_ok=True)
        os.makedirs(mask_folder, exist_ok=True)

    image_files = [
        f for f in os.listdir(img_path)
        if f.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp', '.tif', '.tiff'))
    ]

    # palette_array는 이미지마다 동일하므로 루프 외부에서 1회 생성
    palette_array = np.array(model.PALETTE, dtype=np.uint8)

    inference_cache = {}

    for image_file in tqdm(image_files, desc="Processing images"):
        image_path = os.path.join(img_path, image_file)
        image_name, image_ext = os.path.splitext(image_file)
        image_ext = image_ext.lower()

        # 이미지를 1회만 읽어 inference와 blending 모두에 재사용
        img_bgr = cv2.imread(image_path)
        img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)

        # 이미 로드된 배열을 inference_segmentor에 전달 (경로 전달 시 내부에서 재읽기 발생)
        result = inference_segmentor(model, img_bgr)
        seg_map = result[0]

        # Create color mask from palette
        color_mask = palette_array[seg_map]

        # blended image
        blended = blend_images(img_rgb, color_mask, alpha=0.6)
        blended_bgr = cv2.cvtColor(blended, cv2.COLOR_RGB2BGR)

        # mask — numpy 슬라이싱으로 cv2.cvtColor 호출 1회 제거
        mask_bgr = color_mask[:, :, ::-1].copy()

        # 결과를 메모리 캐시에 저장 (georeference 단계에서 재사용)
        # mask는 palette 원본(RGB) 그대로 저장 — Oilshape의 class_colors가 RGB 기준이므로 BGR로 저장 시 색상 매칭 실패
        inference_cache[image_file] = {
            'blended': blended_bgr,
            'mask': color_mask,
            'ext': image_ext,
        }

        if write_files:
            cv2.imwrite(
                os.path.join(output_folder, f"{image_name}{image_ext}"),
                blended_bgr,
                [cv2.IMWRITE_JPEG_QUALITY, 85]
            )
            cv2.imwrite(os.path.join(mask_folder, f"{image_name}{image_ext}"), mask_bgr)

    return inference_cache


if __name__ == '__main__':
    if len(sys.argv) < 2:
        raise ValueError("파라미터가 제공되지 않았습니다. 폴더 이름을 명령줄 인자로 입력해주세요.")
    _model = load_model()
    # CLI 단독 실행 시에는 중간 파일도 디스크에 저장
    run_inference(_model, sys.argv[1], write_files=True)