import os
import cv2
import numpy as np

DEFAULT_PARAMS = {
    "sobel_kernel": 7,
    "sobel_thresh": (35, 50),
    "r_thresh": (225, 255),
    "s_thresh": (220, 250),
}

def warp_to_birdeye(image, src_points, dst_points):
    # Warp the camera view into bird eye view
    matrix = cv2.getPerspectiveTransform(src_points, dst_points)
    return cv2.warpPerspective(image, matrix, (image.shape[1], image.shape[0]), flags=cv2.INTER_LINEAR)

def warp_back(image, src_points, dst_points):
    # Warp the bird eye view back to the original camera view
    matrix = cv2.getPerspectiveTransform(dst_points, src_points)
    return cv2.warpPerspective(image, matrix, (image.shape[1], image.shape[0]), flags=cv2.INTER_LINEAR)

def sobel_binary(image, mode="xy", kernel=5, thresh=(25, 200)):
    # Compute Sobel gradients and threshold
    gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)

    if mode == "x":
        sobel = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=kernel)
    elif mode == "y":
        sobel = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=kernel)
    else:
        sobel_x = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=kernel)
        sobel_y = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=kernel)
        sobel = np.sqrt(sobel_x**2 + sobel_y**2)

    abs_sobel = np.absolute(sobel)
    max_val = abs_sobel.max()
    if max_val == 0:
        return np.zeros_like(gray, dtype=np.uint8)

    scaled = (255 * abs_sobel / max_val).astype(np.uint8)
    return ((scaled >= thresh[0]) & (scaled <= thresh[1])).astype(np.uint8)

def simple_threshold(channel, thresh):
    # Keep pixels within the range
    return ((channel >= thresh[0]) & (channel <= thresh[1])).astype(np.uint8)

def red_channel_binary(image, thresh):
    # Highlight yellow lane markings
    return simple_threshold(image[:, :, 0], thresh)

def saturation_binary(image, thresh):
    # Highlight bright white lane marking
    hls = cv2.cvtColor(image, cv2.COLOR_RGB2HLS)
    return simple_threshold(hls[:, :, 2], thresh)

def mask_road_region(image):
    # Mask everything outside a trapezoid covering the road
    h, w = image.shape[:2]
    vertices = np.array([[(0, h), (w // 2, h // 2 + 45), (w, h)]], dtype=np.int32)
    mask = np.zeros_like(image)
    fill_value = (255,) * image.shape[2] if image.ndim == 3 else 255
    cv2.fillPoly(mask, vertices, fill_value)
    return cv2.bitwise_and(image, mask)

def build_binary_mask(image, src_points, dst_points, params):
    # Combine color and gradient cues into binary mask
    cropped = mask_road_region(image)
    warped = warp_to_birdeye(cropped, src_points, dst_points)

    red_mask = red_channel_binary(warped, params["r_thresh"])
    saturation_mask = saturation_binary(warped, params["s_thresh"])
    sobel_mask = sobel_binary(
        warped,
        mode="xy",
        kernel=params["sobel_kernel"],
        thresh=params["sobel_thresh"],
    )

    combined = (red_mask | saturation_mask | sobel_mask).astype(np.uint8)
    return combined, warped

def sliding_window_search(binary_mask):
    # Sliding window search.
    h, w = binary_mask.shape
    histogram = np.sum(binary_mask[h // 2 :, :], axis=0)

    midpoint = w // 2
    left_base = np.argmax(histogram[:midpoint])
    right_base = np.argmax(histogram[midpoint:]) + midpoint

    window_height = h // 9
    margin = 60
    minpix = 50

    nonzero_y, nonzero_x = binary_mask.nonzero()
    current_left = left_base
    current_right = right_base
    left_lane_inds = []
    right_lane_inds = []

    for window in range(9):
        win_y_low = h - (window + 1) * window_height
        win_y_high = h - window * window_height
        win_xleft_low = current_left - margin
        win_xleft_high = current_left + margin
        win_xright_low = current_right - margin
        win_xright_high = current_right + margin

        in_window = (nonzero_y >= win_y_low) & (nonzero_y < win_y_high)
        good_left = (in_window & (nonzero_x >= win_xleft_low) & (nonzero_x < win_xleft_high)).nonzero()[0]
        good_right = (in_window & (nonzero_x >= win_xright_low) & (nonzero_x < win_xright_high)).nonzero()[0]

        left_lane_inds.append(good_left)
        right_lane_inds.append(good_right)

        if len(good_left) > minpix:
            current_left = int(nonzero_x[good_left].mean())
        if len(good_right) > minpix:
            current_right = int(nonzero_x[good_right].mean())

    left_lane_inds = np.concatenate(left_lane_inds) if left_lane_inds else np.array([], dtype=np.int64)
    right_lane_inds = np.concatenate(right_lane_inds) if right_lane_inds else np.array([], dtype=np.int64)

    left_fit = np.polyfit(nonzero_y[left_lane_inds], nonzero_x[left_lane_inds], 2) if len(left_lane_inds) else np.array([0, 0, 0])
    right_fit = (
        np.polyfit(nonzero_y[right_lane_inds], nonzero_x[right_lane_inds], 2)
        if len(right_lane_inds)
        else np.array([0, 0, 0])
    )

    vis = np.dstack((binary_mask, binary_mask, binary_mask)) * 255
    vis[nonzero_y[left_lane_inds], nonzero_x[left_lane_inds]] = [255, 0, 0]
    vis[nonzero_y[right_lane_inds], nonzero_x[right_lane_inds]] = [0, 0, 255]

    return left_fit, right_fit, vis

def draw_lane_overlay(original_image, binary_mask, left_fit, right_fit, src_points, dst_points):
    # Draw the detected lane
    h, w = binary_mask.shape
    ploty = np.arange(h, dtype=np.float32)
    ploty_sq = ploty * ploty

    left_x = left_fit[0] * ploty_sq + left_fit[1] * ploty + left_fit[2]
    right_x = right_fit[0] * ploty_sq + right_fit[1] * ploty + right_fit[2]

    lane_canvas = np.zeros((h, w, 3), dtype=np.uint8)
    pts_left = np.column_stack([left_x, ploty]).astype(np.int32)
    pts_right = np.column_stack([right_x[::-1], ploty[::-1]]).astype(np.int32)
    pts = np.vstack([pts_left, pts_right])

    cv2.fillPoly(lane_canvas, [pts], (0, 255, 0))
    cv2.polylines(lane_canvas, [pts_left], False, (255, 0, 255), 10)
    cv2.polylines(lane_canvas, [pts_right], False, (0, 255, 255), 10)

    inverse_warp = warp_back(lane_canvas, src_points, dst_points)
    return cv2.addWeighted(original_image.astype(np.uint8), 1.0, inverse_warp, 0.5, 0)

def default_perspective_points(height, width):
    # Perspective points
    scale_x = width / 1280.0
    scale_y = height / 720.0
    src = np.float32(
        [
            [210 * scale_x, 700 * scale_y],
            [570 * scale_x, 460 * scale_y],
            [705 * scale_x, 460 * scale_y],
            [1075 * scale_x, 700 * scale_y],
        ]
    )
    dst = np.float32(
        [
            [400, height],
            [400, 0],
            [width - 400, 0],
            [width - 400, height],
        ]
    )
    return src, dst

def run_pipeline(image_rgb, params=None):
    params = (params or {}).copy()
    for key, value in DEFAULT_PARAMS.items():
        params.setdefault(key, value)

    height, width = image_rgb.shape[:2]
    src_points, dst_points = default_perspective_points(height, width)

    binary_mask, warped = build_binary_mask(image_rgb, src_points, dst_points, params)
    left_fit, right_fit, search_vis = sliding_window_search(binary_mask)
    final_overlay = draw_lane_overlay(image_rgb, binary_mask, left_fit, right_fit, src_points, dst_points)

    return {
        "binary": binary_mask,
        "warped": warped,
        "sliding_window_vis": search_vis,
        "final": final_overlay,
        "left_fit": left_fit,
        "right_fit": right_fit,
        "src": src_points,
        "dst": dst_points,
        "params": params,
    }

def process_image(image_path, output_dir=None, params=None):
    # Load image, run pipeline
    image_path_str = str(image_path)
    img_bgr = cv2.imread(image_path_str)
    if img_bgr is None:
        raise FileNotFoundError(f"Unable to read image at {image_path_str}")

    img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
    results = run_pipeline(img_rgb, params=params)

    if output_dir is not None:
        output_dir_str = str(output_dir)
        os.makedirs(output_dir_str, exist_ok=True)
        base_name = os.path.splitext(os.path.basename(image_path_str))[0]
        out_path = os.path.join(output_dir_str, f"{base_name}-output.png")
        final_bgr = cv2.cvtColor(results["final"], cv2.COLOR_RGB2BGR)
        cv2.imwrite(out_path, final_bgr)

    return results

__all__ = ["process_image"]