import time
import krpc



def get_surface_speed():
    conn = krpc.connect(name='Vessel speed')
    vessel = conn.space_center.active_vessel
    obt_frame = vessel.orbit.body.non_rotating_reference_frame
    srf_frame = vessel.orbit.body.reference_frame

    while True:
        obt_speed = vessel.flight(obt_frame).speed
        srf_speed = vessel.flight(srf_frame).speed
        print('Orbital speed = %.1f m/s, Surface speed = %.1f m/s' %
              (obt_speed, srf_speed))
        time.sleep(1)

def get_surface_velocity():
    conn = krpc.connect(name='Orbital speed')
    vessel = conn.space_center.active_vessel
    ref_frame = conn.space_center.ReferenceFrame.create_hybrid(
        position=vessel.orbit.body.reference_frame,
        rotation=vessel.surface_reference_frame)

    while True:
        velocity = vessel.flight(ref_frame).velocity
        print('Surface velocity = (%.1f, %.1f, %.1f)' % velocity)
        time.sleep(1)


def surface_prograde():
    conn = krpc.connect(name='Surface prograde')
    vessel = conn.space_center.active_vessel
    ap = vessel.auto_pilot

    ap.reference_frame = vessel.surface_velocity_reference_frame
    ap.target_direction = (0, 1, 0)
    ap.engage()
    ap.wait()
    print('Auto-pilot now pointing prograde in surface reference frame')
    ap.disengage()

def surface_retrograde():
    conn = krpc.connect(name='Surface retrograde')
    vessel = conn.space_center.active_vessel
    ap = vessel.auto_pilot

    ap.reference_frame = vessel.surface_velocity_reference_frame
    ap.target_direction = (0, -1, 0)
    ap.engage()
    ap.wait()
    print('Auto-pilot now pointing retrograde in surface reference frame')
    ap.disengage()

def vessel_available_thrust():
    conn = krpc.connect(name='Vessel available thrust')
    vessel = conn.space_center.active_vessel

    while True:
        thrust = vessel.available_thrust
        print('Vessel available thrust = %.1f kN' % (thrust / 1000))
        time.sleep(1)

def auto_land_predictive(host='localhost', port=50000, target_altitude=0.5, target_vspeed=-2.0, safety_margin=5.0, verbose=True):
    """预测式自动着陆（suicide burn 思路）。

    关键点：
    1) 速度从接口获取：surface speed 用 flight(surface_frame).speed（标量），垂直速度用 flight(surface_frame).vertical_speed。
       不再用“高度差/时间差”估算。
    2) 不再脚本激活即点火；先计算 braking distance，根据可用推力决定何时点火（countdown）。

    近似模型（忽略空气阻力、推力变化、TWR变化导致的积分误差）：
      a_max_up = max(0, available_thrust/mass - g)
      braking_distance ≈ (v^2 - vf^2) / (2*a_max_up) + safety_margin
    当 remaining_alt <= braking_distance 时开始点火。

    燃烧阶段按剩余距离反推所需净向上加速度：
      a_req_up = max(0, (v^2 - vf^2) / (2*remaining_alt))
      throttle = (a_req_up + g) * mass / available_thrust

    注：这里 v 取“相对地表速度”的标量（surface speed），更贴近你提到的 get_surface_speed。
        若需要更精细，可改成取 surface_velocity 向量的径向分量/或用垂直速度替代。
    """
    try:
        import math
    except Exception:
        math = None

    if krpc is None:
        raise RuntimeError('krpc 未安装')

    conn = krpc.connect(name='auto_land_predictive', address=host, rpc_port=port)
    sc = conn.space_center
    vessel = sc.active_vessel
    if vessel is None:
        conn.close()
        raise RuntimeError('未找到活动载具')

    status = {'result': 'failed', 'logs': []}

    def _get_g(local_vessel):
        try:
            body = local_vessel.orbit.body
            g = float(getattr(body, 'surface_gravity', 9.81))
            if g > 0:
                return g
        except Exception:
            pass
        try:
            g = float(getattr(local_vessel.flight(), 'surface_gravity', 9.81))
            if g > 0:
                return g
        except Exception:
            pass
        return 9.81

    # 使用 body 的 reference_frame 读取 surface speed / vertical_speed（对应 get_surface_speed 思路）
    try:
        surface_frame = vessel.orbit.body.reference_frame
    except Exception:
        surface_frame = None

    # autopilot 指向 surface retrograde
    try:
        ap = vessel.auto_pilot
        ap.reference_frame = vessel.surface_velocity_reference_frame
        ap.target_direction = (0, -1, 0)
        ap.engage()
    except Exception:
        ap = None

    def _read_state():
        """从接口读取状态：alt, surface_speed(标量), vertical_speed, mass, available_thrust, g, situation_name"""
        flight = vessel.flight(surface_frame) if surface_frame is not None else vessel.flight()
        alt = float(getattr(flight, 'surface_altitude', 0.0))
        srf_speed = float(getattr(flight, 'speed', 0.0))
        v_speed = float(getattr(flight, 'vertical_speed', 0.0))
        mass = float(getattr(vessel, 'mass', 0.0))
        avail_thrust = float(getattr(vessel, 'available_thrust', 0.0))
        g_local = _get_g(vessel)
        try:
            situation_name = str(getattr(getattr(vessel, 'situation', None), 'name', ''))
        except Exception:
            situation_name = ''
        return alt, srf_speed, v_speed, mass, avail_thrust, g_local, situation_name

    try:
        # 初始确保不点火
        try:
            vessel.control.throttle = 0.0
        except Exception:
            pass

        # 倒计时/等待点火
        ignited = False
        last_log_t = 0.0
        while True:
            alt, srf_speed, v_speed, mass, avail_thrust, g_local, situation_name = _read_state()
            remaining_alt = max(0.0, alt - float(target_altitude))

            # v：使用 surface speed（标量）作为“需要刹掉的速度”
            v = max(0.0, float(srf_speed))
            vf = max(0.0, abs(float(target_vspeed)))

            # 若推力不可用，直接点火尝试（避免永远等不到）
            if mass <= 0 or avail_thrust <= 1e-6:
                a_max_up = 0.0
            else:
                a_max_up = max(0.0, avail_thrust / mass - g_local)

            if a_max_up <= 1e-6:
                braking_dist = float('inf')
            else:
                braking_dist = max(0.0, (v * v - vf * vf) / (2.0 * a_max_up)) + float(safety_margin)

            countdown = None
            if v > 1e-6 and braking_dist != float('inf'):
                # 近似把“还剩多少距离”除以当前速度当作倒计时（粗略）
                countdown = (remaining_alt - braking_dist) / v

            now = time.time()
            if verbose and now - last_log_t > 0.2:
                last_log_t = now
                msg = (
                    f'WAIT alt={alt:.1f} rem={remaining_alt:.1f} '
                    f'srf_speed={srf_speed:.2f} vs={v_speed:.2f} '
                    f'a_max_up={a_max_up:.2f} braking={braking_dist:.1f} '
                    f'countdown={(countdown if countdown is not None else float("nan")):.2f}'
                )
                status['logs'].append(msg)
                print(msg)

            # 点火条件：进入刹车距离 or 推力未知（inf）但高度很低
            should_ignite = (remaining_alt <= braking_dist) or (braking_dist == float('inf') and remaining_alt < 200.0)
            if should_ignite:
                break

            # 确保仍在滑行
            try:
                vessel.control.throttle = 0.0
            except Exception:
                pass
            time.sleep(0.01)

        # 点火
        try:
            vessel.control.activate_next_stage()
        except Exception:
            pass
        ignited = True

        # 燃烧闭环：根据 remaining_alt 反推需要的减速度 -> 油门
        while True:
            alt, srf_speed, v_speed, mass, avail_thrust, g_local, situation_name = _read_state()
            remaining_alt = max(0.0, alt - float(target_altitude))

            v = max(0.0, float(srf_speed))
            vf = max(0.0, abs(float(target_vspeed)))

            # 终止条件：以 vessel.situation.name == 'landed' 为准
            if str(situation_name).strip().lower() == 'landed':
                try:
                    vessel.control.throttle = 0.0
                except Exception:
                    pass
                status['result'] = 'landed'
                status['logs'].append('检测到 vessel.situation=LANDED，着陆完成')
                if verbose:
                    print('检测到 vessel.situation=LANDED，着陆完成')
                break

            # 兜底：极低高度 + 速度足够小（防止 situation 迟迟不切换导致卡死）
            if remaining_alt <= max(0.2, float(target_altitude)) and v <= (vf + 0.5):
                try:
                    vessel.control.throttle = 0.0
                except Exception:
                    pass
                status['result'] = 'touchdown_pending'
                status['logs'].append('达到高度/速度阈值，但 vessel.situation 仍未变为 LANDED')
                if verbose:
                    print('达到高度/速度阈值，但 vessel.situation 仍未变为 LANDED')
                break

            # 计算所需净向上加速度
            if remaining_alt <= 1e-3:
                a_req_up = 0.0
            else:
                a_req_up = max(0.0, (v * v - vf * vf) / (2.0 * remaining_alt))

            # 推力不足时给满油门
            if mass <= 0 or avail_thrust <= 1e-6:
                throttle = 1.0
            else:
                throttle = (a_req_up + g_local) * mass / avail_thrust

            throttle = max(0.0, min(1.0, float(throttle)))

            # 低空兜底：高度<5m 且仍很快 -> 满油门
            if alt < 5.0 and v > 10.0:
                throttle = 1.0

            try:
                vessel.control.throttle = throttle
            except Exception:
                pass

            if verbose:
                msg = (
                    f'BURN alt={alt:.1f} rem={remaining_alt:.1f} '
                    f'srf_speed={srf_speed:.2f} vs={v_speed:.2f} '
                    f'a_req_up={a_req_up:.2f} throttle={throttle:.2f}'
                )
                status['logs'].append(msg)
                print(msg)

            time.sleep(0.01)

    finally:
        try:
            if ap is not None and hasattr(ap, 'disengage'):
                try:
                    ap.disengage()
                except Exception:
                    pass
        except Exception:
            pass
        try:
            conn.close()
        except Exception:
            pass

    return status

if __name__ == "__main__":
    xxx = auto_land_predictive()
    print(xxx)
    print('done')