"""Wireframe rendering effect."""

import logging
from typing import Tuple, Optional

import torch

logger = logging.getLogger(__name__)


def render_wireframe(
    vertices: torch.Tensor,
    faces: torch.Tensor,
    mvp: torch.Tensor,
    width: int,
    height: int,
    line_width: float = 1.0,
    line_color: Tuple[float, float, float] = (1.0, 1.0, 1.0),
    background_color: Tuple[float, float, float] = (0.0, 0.0, 0.0),
    fill_color: Optional[Tuple[float, float, float]] = None,
) -> torch.Tensor:
    """
    Render a mesh as wireframe.

    Args:
        vertices: Mesh vertices [N, 3]
        faces: Face indices [F, 3]
        mvp: Model-view-projection matrix [4, 4]
        width: Output width
        height: Output height
        line_width: Width of wireframe lines
        line_color: RGB color for lines
        background_color: RGB background color
        fill_color: Optional fill color for faces

    Returns:
        Rendered wireframe image [H, W, 3]
    """
    device = vertices.device

    # Transform vertices to clip space
    verts_homo = torch.cat([
        vertices,
        torch.ones(vertices.shape[0], 1, device=device)
    ], dim=1)
    verts_clip = (mvp @ verts_homo.T).T

    # Perspective divide to NDC
    verts_ndc = verts_clip[:, :3] / (verts_clip[:, 3:4] + 1e-8)

    # Convert to screen space
    verts_screen = torch.zeros_like(verts_ndc[:, :2])
    verts_screen[:, 0] = (verts_ndc[:, 0] + 1) * 0.5 * width
    verts_screen[:, 1] = (1 - verts_ndc[:, 1]) * 0.5 * height  # Flip Y

    # Initialize output
    output = torch.zeros(height, width, 3, dtype=torch.float32, device=device)
    output[:, :, 0] = background_color[0]
    output[:, :, 1] = background_color[1]
    output[:, :, 2] = background_color[2]

    # Draw edges
    line_color_t = torch.tensor(line_color, dtype=torch.float32, device=device)

    for face in faces:
        for j in range(3):
            v0 = verts_screen[face[j]]
            v1 = verts_screen[face[(j + 1) % 3]]

            # Draw line between v0 and v1
            output = _draw_line(output, v0, v1, line_color_t, line_width)

    return output


def wireframe_overlay(
    base_image: torch.Tensor,
    wireframe: torch.Tensor,
    opacity: float = 0.5,
    blend_mode: str = "over"
) -> torch.Tensor:
    """
    Overlay wireframe on top of base image.

    Args:
        base_image: Base rendered image [H, W, 3]
        wireframe: Wireframe image [H, W, 3]
        opacity: Wireframe opacity
        blend_mode: "over", "add", or "screen"

    Returns:
        Composited image [H, W, 3]
    """
    if blend_mode == "add":
        return (base_image + wireframe * opacity).clamp(0, 1)
    elif blend_mode == "screen":
        return (1 - (1 - base_image) * (1 - wireframe * opacity)).clamp(0, 1)
    else:  # over
        # Only overlay where wireframe has color
        mask = (wireframe.sum(dim=-1, keepdim=True) > 0).float()
        return base_image * (1 - mask * opacity) + wireframe * mask * opacity


def _draw_line(
    image: torch.Tensor,
    p0: torch.Tensor,
    p1: torch.Tensor,
    color: torch.Tensor,
    width: float = 1.0
) -> torch.Tensor:
    """
    Draw a line on an image using Bresenham's algorithm.

    Args:
        image: Image tensor [H, W, 3]
        p0: Start point [2]
        p1: End point [2]
        color: Line color [3]
        width: Line width

    Returns:
        Image with line drawn
    """
    H, W, _ = image.shape

    x0, y0 = int(p0[0].item()), int(p0[1].item())
    x1, y1 = int(p1[0].item()), int(p1[1].item())

    # Bresenham's line algorithm
    dx = abs(x1 - x0)
    dy = abs(y1 - y0)
    sx = 1 if x0 < x1 else -1
    sy = 1 if y0 < y1 else -1
    err = dx - dy

    half_w = int(width / 2)

    while True:
        # Draw pixel with width
        for wx in range(-half_w, half_w + 1):
            for wy in range(-half_w, half_w + 1):
                px, py = x0 + wx, y0 + wy
                if 0 <= px < W and 0 <= py < H:
                    image[py, px] = color

        if x0 == x1 and y0 == y1:
            break

        e2 = 2 * err
        if e2 > -dy:
            err -= dy
            x0 += sx
        if e2 < dx:
            err += dx
            y0 += sy

    return image


def extract_edges(
    faces: torch.Tensor
) -> torch.Tensor:
    """
    Extract unique edges from face indices.

    Args:
        faces: Face indices [F, 3]

    Returns:
        Edge indices [E, 2]
    """
    # Get all edges (including duplicates)
    edges = []
    for i in range(3):
        j = (i + 1) % 3
        e = torch.stack([faces[:, i], faces[:, j]], dim=1)
        edges.append(e)

    all_edges = torch.cat(edges, dim=0)

    # Sort each edge so smaller index comes first
    sorted_edges = torch.sort(all_edges, dim=1).values

    # Remove duplicates
    unique_edges = torch.unique(sorted_edges, dim=0)

    return unique_edges


def compute_edge_visibility(
    vertices: torch.Tensor,
    faces: torch.Tensor,
    viewmat: torch.Tensor,
    edges: Optional[torch.Tensor] = None
) -> torch.Tensor:
    """
    Compute which edges are visible (silhouette edges).

    An edge is a silhouette edge if its two adjacent faces have
    different front/back facing orientations.

    Args:
        vertices: Vertex positions [N, 3]
        faces: Face indices [F, 3]
        viewmat: View matrix [4, 4]
        edges: Optional pre-computed edges [E, 2]

    Returns:
        Boolean mask of visible edges [E]
    """
    device = vertices.device

    if edges is None:
        edges = extract_edges(faces)

    # Compute face normals
    v0 = vertices[faces[:, 0]]
    v1 = vertices[faces[:, 1]]
    v2 = vertices[faces[:, 2]]

    face_normals = torch.cross(v1 - v0, v2 - v0)
    face_normals = face_normals / (torch.norm(face_normals, dim=1, keepdim=True) + 1e-8)

    # Face centers
    face_centers = (v0 + v1 + v2) / 3

    # View direction (from face center to camera)
    cam_pos = -viewmat[:3, :3].T @ viewmat[:3, 3]
    view_dirs = cam_pos.unsqueeze(0) - face_centers
    view_dirs = view_dirs / (torch.norm(view_dirs, dim=1, keepdim=True) + 1e-8)

    # Front-facing test
    dots = (face_normals * view_dirs).sum(dim=1)
    front_facing = dots > 0

    # Build edge-to-face mapping (simplified: mark edge visible if any adjacent face is front-facing)
    # For proper silhouette, need to check if edge is between front and back facing faces
    visible = torch.ones(len(edges), dtype=torch.bool, device=device)

    return visible