import os
import tempfile
from dataclasses import dataclass
from typing import List, Tuple

import gradio as gr
import imageio.v3 as iio
import numpy as np

try:
    import cv2  # type: ignore
except ImportError:  # pragma: no cover - optional dependency
    cv2 = None


SUPPORTED_EXTENSIONS = {".png", ".tif", ".tiff", ".exr"}


@dataclass
class MeshData:
    vertices: List[Tuple[float, float, float]]
    faces: List[Tuple[int, int, int]]


def _normalize_heightmap(heightmap: np.ndarray) -> np.ndarray:
    if heightmap.ndim == 3:
        if heightmap.shape[2] >= 3:
            heightmap = heightmap[..., :3].mean(axis=2)
        else:
            heightmap = heightmap[..., 0]

    heightmap = np.asarray(heightmap, dtype=np.float32)
    min_val = float(np.min(heightmap))
    max_val = float(np.max(heightmap))
    if max_val <= min_val:
        return np.zeros_like(heightmap, dtype=np.float32)
    return (heightmap - min_val) / (max_val - min_val)


def _read_heightmap(path: str) -> Tuple[np.ndarray, bool]:
    ext = os.path.splitext(path)[1].lower()
    if ext not in SUPPORTED_EXTENSIONS:
        raise ValueError(f"Unsupported file type: {ext}")

    if ext == ".exr":
        if cv2 is None:
            raise RuntimeError("EXR読み込みにはopencv-pythonが必要です。")
        image = cv2.imread(path, cv2.IMREAD_ANYDEPTH | cv2.IMREAD_GRAYSCALE)
        if image is None:
            raise RuntimeError("EXRファイルの読み込みに失敗しました。")
        return _normalize_heightmap(image), False

    image = iio.imread(path)
    is_8bit = image.dtype == np.uint8
    return _normalize_heightmap(image), is_8bit


def _smooth_heightmap(heightmap: np.ndarray) -> np.ndarray:
    if heightmap.ndim != 2:
        raise ValueError("高さマップは2次元配列である必要があります。")
    padded = np.pad(heightmap, 1, mode="edge")
    blurred = (
        padded[:-2, :-2]
        + 2 * padded[:-2, 1:-1]
        + padded[:-2, 2:]
        + 2 * padded[1:-1, :-2]
        + 4 * padded[1:-1, 1:-1]
        + 2 * padded[1:-1, 2:]
        + padded[2:, :-2]
        + 2 * padded[2:, 1:-1]
        + padded[2:, 2:]
    ) / 16.0
    return blurred.astype(np.float32, copy=False)


def _build_mesh(
    heightmap: np.ndarray,
    x_mm_per_px: float,
    y_mm_per_px: float,
    base_thickness_mm: float,
    height_scale_mm: float,
    bottom_mode: str,
    reverse_face: bool,
) -> MeshData:
    heightmap = np.clip(heightmap, 0.0, 1.0)
    height_values = base_thickness_mm + heightmap * height_scale_mm

    rows, cols = heightmap.shape
    vertices: List[Tuple[float, float, float]] = []
    faces: List[Tuple[int, int, int]] = []

    for r in range(rows):
        y = (rows - 1 - r) * y_mm_per_px
        for c in range(cols):
            x = c * x_mm_per_px
            z = float(height_values[r, c])
            vertices.append((x, y, z))

    top_offset = 0
    for r in range(rows - 1):
        for c in range(cols - 1):
            v0 = top_offset + r * cols + c
            v1 = top_offset + r * cols + c + 1
            v2 = top_offset + (r + 1) * cols + c
            v3 = top_offset + (r + 1) * cols + c + 1
            faces.append((v0, v1, v2))
            faces.append((v1, v3, v2))

    def perimeter_indices() -> List[int]:
        if rows == 1 and cols == 1:
            return [0]
        indices = [c for c in range(cols)]
        if rows > 2:
            indices += [r * cols + (cols - 1) for r in range(1, rows - 1)]
        if rows > 1:
            indices += [(rows - 1) * cols + c for c in range(cols - 1, -1, -1)]
        if cols > 1 and rows > 2:
            indices += [r * cols for r in range(rows - 2, 0, -1)]
        return indices

    bottom_offset = len(vertices)
    bottom_index_map = None
    if bottom_mode == "normal":
        for r in range(rows):
            y = (rows - 1 - r) * y_mm_per_px
            for c in range(cols):
                x = c * x_mm_per_px
                z = 0.0
                vertices.append((x, y, z))

        for r in range(rows - 1):
            for c in range(cols - 1):
                v0 = bottom_offset + r * cols + c
                v1 = bottom_offset + r * cols + c + 1
                v2 = bottom_offset + (r + 1) * cols + c
                v3 = bottom_offset + (r + 1) * cols + c + 1
                faces.append((v0, v2, v1))
                faces.append((v1, v2, v3))
    else:
        perimeter = perimeter_indices()
        bottom_index_map = {}
        for top_idx in perimeter:
            r = top_idx // cols
            c = top_idx % cols
            x = c * x_mm_per_px
            y = (rows - 1 - r) * y_mm_per_px
            z = 0.0
            bottom_index_map[top_idx] = len(vertices)
            vertices.append((x, y, z))

        if bottom_mode == "reduction" and len(perimeter) >= 3:
            center = bottom_index_map[perimeter[0]]
            for i in range(1, len(perimeter) - 1):
                a = bottom_index_map[perimeter[i]]
                b = bottom_index_map[perimeter[i + 1]]
                faces.append((center, b, a))

    def side_faces(index_iter, flip=False):
        for i in range(len(index_iter) - 1):
            top_a = index_iter[i]
            top_b = index_iter[i + 1]
            if bottom_index_map is None:
                bottom_a = bottom_offset + top_a
                bottom_b = bottom_offset + top_b
            else:
                bottom_a = bottom_index_map[top_a]
                bottom_b = bottom_index_map[top_b]
            if flip:
                faces.append((top_a, bottom_b, bottom_a))
                faces.append((top_a, top_b, bottom_b))
            else:
                faces.append((top_a, bottom_a, bottom_b))
                faces.append((top_a, bottom_b, top_b))

    # Top edge (row 0)
    side_faces([c for c in range(cols)], flip=False)
    # Bottom edge (row rows-1)
    side_faces([ (rows - 1) * cols + c for c in range(cols)], flip=True)
    # Left edge (col 0)
    side_faces([ r * cols for r in range(rows)], flip=True)
    # Right edge (col cols-1)
    side_faces([ r * cols + (cols - 1) for r in range(rows)], flip=False)

    if reverse_face:
        return MeshData(vertices=vertices, faces=faces)
    flipped_faces = [(a, c, b) for a, b, c in faces]
    return MeshData(vertices=vertices, faces=flipped_faces)


def _resample_heightmap(heightmap: np.ndarray, factor: float) -> np.ndarray:
    if factor <= 0:
        raise ValueError("XY解像度係数は正の値である必要があります。")
    if abs(factor - 1.0) < 1e-6:
        return heightmap

    rows, cols = heightmap.shape
    new_rows = max(1, int(round(rows * factor)))
    new_cols = max(1, int(round(cols * factor)))

    if cv2 is not None:
        interpolation = cv2.INTER_AREA if factor < 1.0 else cv2.INTER_LINEAR
        resized = cv2.resize(heightmap, (new_cols, new_rows), interpolation=interpolation)
        return resized.astype(np.float32, copy=False)

    row_coords = np.linspace(0, rows - 1, new_rows)
    col_coords = np.linspace(0, cols - 1, new_cols)
    tmp = np.empty((new_rows, cols), dtype=np.float32)
    for idx, coord in enumerate(row_coords):
        base = int(np.floor(coord))
        frac = coord - base
        if base >= rows - 1:
            tmp[idx, :] = heightmap[-1, :]
        else:
            tmp[idx, :] = heightmap[base, :] * (1.0 - frac) + heightmap[base + 1, :] * frac

    resized = np.empty((new_rows, new_cols), dtype=np.float32)
    for idx, coord in enumerate(col_coords):
        base = int(np.floor(coord))
        frac = coord - base
        if base >= cols - 1:
            resized[:, idx] = tmp[:, -1]
        else:
            resized[:, idx] = tmp[:, base] * (1.0 - frac) + tmp[:, base + 1] * frac

    return resized


def _write_obj(mesh: MeshData, out_path: str) -> str:
    with open(out_path, "w", encoding="utf-8") as obj_file:
        obj_file.write("# Heightmap mesh\n")
        for v in mesh.vertices:
            obj_file.write(f"v {v[0]:.6f} {v[1]:.6f} {v[2]:.6f}\n")
        for f in mesh.faces:
            v1, v2, v3 = (idx + 1 for idx in f)
            obj_file.write(f"f {v1} {v2} {v3}\n")
    return out_path


def _write_ply(mesh: MeshData, out_path: str) -> str:
    with open(out_path, "w", encoding="utf-8") as ply_file:
        ply_file.write("ply\nformat ascii 1.0\n")
        ply_file.write(f"element vertex {len(mesh.vertices)}\n")
        ply_file.write("property float x\nproperty float y\nproperty float z\n")
        ply_file.write(f"element face {len(mesh.faces)}\n")
        ply_file.write("property list uchar int vertex_indices\nend_header\n")
        for v in mesh.vertices:
            ply_file.write(f"{v[0]:.6f} {v[1]:.6f} {v[2]:.6f}\n")
        for f in mesh.faces:
            ply_file.write(f"3 {f[0]} {f[1]} {f[2]}\n")
    return out_path


def generate_mesh(
    file,
    output_format,
    height_ratio,
    xy_resolution_factor,
    width_mm,
    height_mm,
    base_thickness_mm,
    bottom_mode,
    invert_heightmap,
    reverse_face,
):
    if file is None:
        raise gr.Error("ファイルをアップロードしてください。")

    if width_mm <= 0 or height_mm <= 0:
        raise gr.Error("縦横サイズ(mm)は正の値で指定してください。")

    path = file.name if hasattr(file, "name") else file
    heightmap, is_8bit = _read_heightmap(path)
    if is_8bit:
        heightmap = _smooth_heightmap(heightmap)
    heightmap = _resample_heightmap(heightmap, float(xy_resolution_factor))
    if invert_heightmap:
        heightmap = 1.0 - heightmap

    height_scale_mm = float(height_ratio) * float(np.sqrt(width_mm * height_mm))
    bottom_mode_map = {
        "ノーマル": "normal",
        "リダクション": "reduction",
        "カット": "cut",
    }
    resolved_bottom_mode = bottom_mode_map.get(str(bottom_mode), "reduction")

    rows, cols = heightmap.shape
    x_mm_per_px = float(width_mm) / max(cols - 1, 1)
    y_mm_per_px = float(height_mm) / max(rows - 1, 1)

    mesh = _build_mesh(
        heightmap=heightmap,
        x_mm_per_px=x_mm_per_px,
        y_mm_per_px=y_mm_per_px,
        base_thickness_mm=float(base_thickness_mm),
        height_scale_mm=float(height_scale_mm),
        bottom_mode=resolved_bottom_mode,
        reverse_face=bool(reverse_face),
    )

    suffix = ".obj" if output_format == "OBJ" else ".ply"
    with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp_file:
        out_path = tmp_file.name

    if output_format == "OBJ":
        _write_obj(mesh, out_path)
    else:
        _write_ply(mesh, out_path)

    return out_path


def preview_heightmap(file):
    if file is None:
        return None
    path = file.name if hasattr(file, "name") else file
    heightmap, _ = _read_heightmap(path)
    return heightmap


def _update_aspect_from_file(file, width_mm, height_mm, lock_aspect):
    if file is None:
        return 1.0, gr.update(), gr.update()
    path = file.name if hasattr(file, "name") else file
    heightmap, _ = _read_heightmap(path)
    rows, cols = heightmap.shape
    aspect = cols / rows if rows else 1.0
    if lock_aspect:
        return aspect, gr.update(value=float(width_mm)), gr.update(value=float(width_mm) / aspect)
    return aspect, gr.update(), gr.update()


def _sync_height_from_width(width_mm, aspect, lock_aspect, sync_source):
    if sync_source == "height":
        return gr.update(), ""
    if not lock_aspect or aspect <= 0:
        return gr.update(), ""
    return gr.update(value=float(width_mm) / aspect), "width"


def _sync_width_from_height(height_mm, aspect, lock_aspect, sync_source):
    if sync_source == "width":
        return gr.update(), ""
    if not lock_aspect or aspect <= 0:
        return gr.update(), ""
    return gr.update(value=float(height_mm) * aspect), "height"


def _height_scale_from_ratio(height_ratio, width_mm, height_mm):
    return float(height_ratio) * float(np.sqrt(float(width_mm) * float(height_mm)))


def _height_ratio_from_scale(height_scale_mm, width_mm, height_mm):
    area = float(width_mm) * float(height_mm)
    if area <= 0:
        return 0.0
    return float(height_scale_mm) / float(np.sqrt(area))


def _sync_height_scale_from_ratio(height_ratio, width_mm, height_mm, height_source):
    if height_source == "scale":
        return gr.update(), ""
    return gr.update(value=_height_scale_from_ratio(height_ratio, width_mm, height_mm)), "ratio"


def _sync_ratio_from_height_scale(
    height_scale_mm, width_mm, height_mm, height_source
):
    if height_source == "ratio":
        return gr.update(), ""
    height_ratio = _height_ratio_from_scale(height_scale_mm, width_mm, height_mm)
    return gr.update(value=height_ratio), "scale"


def _sync_height_inputs_from_size(
    width_mm, height_mm, height_ratio, height_scale_mm, height_source
):
    if height_source == "scale":
        height_ratio = _height_ratio_from_scale(height_scale_mm, width_mm, height_mm)
        return gr.update(value=height_ratio), gr.update()
    return (
        gr.update(),
        gr.update(value=_height_scale_from_ratio(height_ratio, width_mm, height_mm)),
    )


def build_app() -> gr.Blocks:
    with gr.Blocks(
        title="Heightmap to Mesh",
        css="""
        #preview-image button[aria-label="Share"] {
            display: none !important;
        }
        #preview-image button[title="Share"],
        #preview-image .share-btn,
        #preview-image .share-button {
            display: none !important;
        }
        """,
    ) as demo:
        gr.Markdown(
            """
            # Heightmap to Mesh (OBJ/PLY)
            モノクロのハイトマップ画像からOBJ/PLYメッシュを生成します。
            - 対応形式: PNG / TIFF / EXR - 8bit/16bit/float
            """
        )
        with gr.Row():
            with gr.Column():
                file_input = gr.File(label="ハイトマップ画像")
                with gr.Row():
                    reverse_face = gr.Checkbox(value=False, label="ポリゴン裏表反転")
                    output_format = gr.Radio(["OBJ", "PLY"], value="OBJ", label="出力形式")
                xy_resolution_input = gr.Number(
                    value=0.5,
                    minimum=0.001,
                    label="XY解像度",
                )
                bottom_mode = gr.Dropdown(
                    ["ノーマル", "リダクション", "カット"],
                    value="リダクション",
                    label="底面メッシュ",
                )
                with gr.Row():
                    lock_aspect = gr.Checkbox(value=True, label="比率固定")
                    width_mm = gr.Number(value=100.0, minimum=0.001, label="横サイズ(mm)")
                    height_mm = gr.Number(value=100.0, minimum=0.001, label="縦サイズ(mm)")
                with gr.Row():
                    invert_heightmap = gr.Checkbox(value=False, label="ハイトマップ上下反転")
                    height_ratio_input = gr.Number(
                        value=0.5,
                        minimum=0.0,
                        label="高さ比率",
                    )
                    height_scale_mm = gr.Number(
                        value=_height_scale_from_ratio(0.5, 100.0, 100.0),
                        minimum=0.0,
                        label="高さ(mm)",
                    )
                base_thickness = gr.Number(value=1.0, minimum=0.0, label="ベース厚み(mm)")
                generate_button = gr.Button("メッシュ生成")
            with gr.Column():
                preview_image = gr.Image(label="プレビュー", type="numpy", elem_id="preview-image")
                output_file = gr.File(label="ダウンロード")
        aspect_state = gr.State(1.0)
        sync_state = gr.State("")
        height_source = gr.State("ratio")
        generate_button.click(
            generate_mesh,
            inputs=[
                file_input,
                output_format,
                height_ratio_input,
                xy_resolution_input,
                width_mm,
                height_mm,
                base_thickness,
                bottom_mode,
                invert_heightmap,
                reverse_face,
            ],
            outputs=output_file,
        )
        file_input.change(preview_heightmap, inputs=file_input, outputs=preview_image)
        file_input.change(
            _update_aspect_from_file,
            inputs=[file_input, width_mm, height_mm, lock_aspect],
            outputs=[aspect_state, width_mm, height_mm],
        )
        width_mm.change(
            _sync_height_from_width,
            inputs=[width_mm, aspect_state, lock_aspect, sync_state],
            outputs=[height_mm, sync_state],
        )
        width_mm.change(
            _sync_height_inputs_from_size,
            inputs=[
                width_mm,
                height_mm,
                height_ratio_input,
                height_scale_mm,
                height_source,
            ],
            outputs=[height_ratio_input, height_scale_mm],
        )
        height_mm.change(
            _sync_width_from_height,
            inputs=[height_mm, aspect_state, lock_aspect, sync_state],
            outputs=[width_mm, sync_state],
        )
        height_mm.change(
            _sync_height_inputs_from_size,
            inputs=[
                width_mm,
                height_mm,
                height_ratio_input,
                height_scale_mm,
                height_source,
            ],
            outputs=[height_ratio_input, height_scale_mm],
        )
        height_ratio_input.change(
            _sync_height_scale_from_ratio,
            inputs=[height_ratio_input, width_mm, height_mm, height_source],
            outputs=[height_scale_mm, height_source],
        )
        height_scale_mm.change(
            _sync_ratio_from_height_scale,
            inputs=[height_scale_mm, width_mm, height_mm, height_source],
            outputs=[height_ratio_input, height_source],
        )
        lock_aspect.change(
            _sync_height_from_width,
            inputs=[width_mm, aspect_state, lock_aspect, sync_state],
            outputs=[height_mm, sync_state],
        )
    return demo


app = build_app()

if __name__ == "__main__":
    port_value = os.getenv("PORT")
    app.launch(
        server_name="0.0.0.0",
        server_port=int(port_value) if port_value else 7860,
    )