import cv2
import os
import sys
import time
import argparse

# 添加父目录到系统路径
current_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.dirname(os.path.dirname(current_dir))
sys.path.append(project_root)

from utils.detect_track import detect_horizontal_track_edge_v3

def process_image(image_path, save_dir=None, show_steps=False):
    """处理单张图像"""
    print(f"处理图像: {image_path}")
    
    # 检测赛道并估算距离
    start_time = time.time()
    edge_point, edge_info = detect_horizontal_track_edge_v3(image_path, observe=show_steps, save_log=True, delay=800)
    processing_time = time.time() - start_time
    
    # 输出结果
    if edge_point is not None and edge_info is not None:
        print(f"处理时间: {processing_time:.3f}秒")
        print(f"边缘点: ({edge_point[0]}, {edge_point[1]})")
        print(f"到中线距离: {edge_info['distance_to_center']}像素")
        print(f"边缘斜率: {edge_info['slope']:.4f}")
        print(f"是否水平: {edge_info['is_horizontal']}")
        print(f"点数量: {edge_info['points_count']}")
        print(f"中线交点: ({edge_info['intersection_point'][0]}, {edge_info['intersection_point'][1]})")
        print(f"交点到底部距离: {edge_info['distance_to_bottom']:.1f}像素")
        print(f"注意: 中线交点是垂直中线与边缘横线的交点")
        print("-" * 30)
        
        # 如果指定了保存目录，保存结果
        if save_dir:
            if not os.path.exists(save_dir):
                os.makedirs(save_dir)
    else:
        print("未能检测到黄色赛道")
    
    return edge_point, edge_info

def main():
    parser = argparse.ArgumentParser(description='黄色赛道检测演示程序')
    parser.add_argument('--input', type=str, default='res/path/task-5/origin_horizontal_edge_20250528_100604_158105.jpg', help='输入图像或视频的路径')
    parser.add_argument('--output', type=str, default='res/path/test-v2/2-end.jpg', help='输出结果的保存路径')
    parser.add_argument('--type', type=str, choices=['image', 'video'], help='输入类型，不指定会自动检测')
    parser.add_argument('--show', default=True, action='store_true', help='显示处理步骤')
    
    args = parser.parse_args()
    
    # 检查输入路径
    if not os.path.exists(args.input):
        print(f"错误：文件 '{args.input}' 不存在")
        return
    
    # 如果未指定类型，根据文件扩展名判断
    if args.type is None:
        ext = os.path.splitext(args.input)[1].lower()
        if ext in ['.jpg', '.jpeg', '.png', '.bmp']:
            args.type = 'image'
        elif ext in ['.mp4', '.avi', '.mov']:
            args.type = 'video'
        else:
            print(f"错误：无法确定文件类型 '{ext}'")
            return
    
    # 根据类型处理
    if args.type == 'image':
        # 获取输出目录
        output_dir = os.path.dirname(args.output)
        process_image(args.input, output_dir, args.show)
    else:
        print("错误：不支持的视频类型")

if __name__ == "__main__":
    main()