"""
KRPC integration helpers
提供一个函数 run_flight_sequence 来连接本地 kRPC 服务，并按照用户需求控制当前活动载具。
需要安装 krpc: pip install krpc

行为：
1) 点火并将油门设为 70%，垂直向上飞行 40s
2) 计算并设定维持 1G 的油门（若计算失败回退到 50%），航向 90，俯仰 45°，维持 20s
3) 将油门提升到 100%，航向 270，俯仰 70°，维持 15s

注意：不同飞行环境（行星/轨道）下重力与最大推力不同，计算油门只是近似。
"""

import time

try:
    import krpc
except Exception:  # 保持导入失败时的可读性
    krpc = None


def run_flight_sequence(host='localhost', port=50000, ignite_stage=True, verbose=True):
    """连接 kRPC 并执行飞行序列。

    参数:
    - host, port: krpc 服务器地址和端口
    - ignite_stage: 是否在开始时执行一次 stage 激活以点火（默认 True）
    - verbose: 是否打印运行信息

    返回: dict 状态信息
    """
    if krpc is None:
        raise RuntimeError('krpc 库未安装或不可用，运行前请 pip install krpc')

    conn = krpc.connect(name='general_tools_flight_sequence', address=host, rpc_port=port)
    sc = conn.space_center

    vessel = sc.active_vessel
    if vessel is None:
        conn.close()
        raise RuntimeError('未找到活动载具')

    # 预先包含 result 字段，使用字符串默认值以保持类型一致
    status = {'steps': [], 'result': ''}

    try:
        # 保险起见先把油门设为 0
        vessel.control.throttle = 0.0
        time.sleep(0.1)

        if ignite_stage:
            if verbose:
                print('激活第一个 stage 以点火（如适用）')
            try:
                # 在分级前先设置目标油门，这样发动机激活后会立刻使用该油门
                vessel.control.throttle = 0.7
                vessel.control.activate_next_stage()
                # activate_next_stage 可能改变当前活动载具（比如 decouple 导致控制权转移），
                # 重新获取 active_vessel 保证后续写入作用在当前载具上
                time.sleep(0.05)
                vessel = sc.active_vessel
            except Exception:
                # 某些情况下 activate_next_stage 可能无效，继续也没问题
                if verbose:
                    print('activate_next_stage 失败，继续执行')

            # 使用 streams 优先监控并确保油门生效
            thrust_ok = _ensure_throttle_applied(conn, sc, desired_throttle=0.7, timeout=10.0, verbose=verbose)
            if not thrust_ok:
                if verbose:
                    print('警告：通过 streams/轮询仍未检测到油门生效，尝试更激进的写入（_force_apply_throttle）并输出发动机诊断')
                try:
                    forced = _force_apply_throttle(conn, sc, desired_throttle=0.7, timeout=5.0, verbose=verbose)
                except Exception:
                    forced = False
                if verbose:
                    print(f'_force_apply_throttle 返回: {forced}')
                try:
                    print(_diagnose_engines(sc.active_vessel))
                except Exception:
                    pass
            time.sleep(0.15)

        # 1) 推力 70%，垂直向上 40s
        t0 = time.time()
        vessel.control.throttle = 0.7
        # 设置自动驾驶为垂直向上（俯仰 90°），保持当前航向
        flight = vessel.flight()
        heading = flight.heading if hasattr(flight, 'heading') else 0.0
        ap = vessel.auto_pilot
        try:
            ap.engage()
            ap.target_pitch_and_heading(90.0, heading)
        except Exception:
            # 如果 target_pitch_and_heading 不可用，尝试设置朝向为航向不变，俯仰 90
            if verbose:
                print('auto_pilot.target_pitch_and_heading 不可用，尝试基础设置')
        if verbose:
            print('阶段1: 70% 推力，垂直上升，持续 40s')
        status['steps'].append({'step': 1, 'throttle': 0.7, 'pitch': 90, 'heading': heading, 'duration_s': 40})
        _sleep_with_progress(40)

        # 2) 计算维持 1G 所需油门，航向 90，俯仰 45，持续 20s
        desired_g = 9.81
        throttle_for_1g = None
        try:
            # vessel.mass 与 vessel.available_thrust 在多数 kRPC 版本可用
            mass = vessel.mass
            available_thrust = vessel.available_thrust
            if available_thrust is None or available_thrust == 0:
                raise RuntimeError('available_thrust 无效')
            # 所需油门比例 = (质量 * 目标加速度) / 可用推力
            throttle_for_1g = (mass * desired_g) / available_thrust
            throttle_for_1g = max(0.0, min(1.0, throttle_for_1g))
        except Exception:
            # 计算失败时回退到一个保守值
            throttle_for_1g = 0.5
            if verbose:
                print('无法精确计算 1G 所需油门，回退到 50%')

        vessel.control.throttle = float(throttle_for_1g)
        try:
            ap.target_pitch_and_heading(45.0, 90.0)
        except Exception:
            if verbose:
                print('设置 45°/90° 方向到 auto_pilot 失败')
        if verbose:
            print(f'阶段2: 维持 1G (估计油门 {throttle_for_1g:.3f})，航向 90，俯仰 45，持续 20s')
        status['steps'].append({'step': 2, 'throttle': throttle_for_1g, 'pitch': 45, 'heading': 90, 'duration_s': 20})
        _sleep_with_progress(20)

        # 3) 推力 100%，航向 270，俯仰 70，持续 15s
        vessel.control.throttle = 1.0
        try:
            ap.target_pitch_and_heading(70.0, 270.0)
        except Exception:
            if verbose:
                print('设置 70°/270° 方向到 auto_pilot 失败')
        if verbose:
            print('阶段3: 100% 推力，航向 270，俯仰 70，持续 15s')
        status['steps'].append({'step': 3, 'throttle': 1.0, 'pitch': 70, 'heading': 270, 'duration_s': 15})
        _sleep_with_progress(15)

        if verbose:
            print('飞行序列完成')
        status['result'] = 'ok'
    finally:
        try:
            # 结束时不改变油门状态，仅可选地切换到 0
            vessel.control.throttle = 0.0
            if hasattr(vessel.auto_pilot, 'disengage'):
                try:
                    vessel.auto_pilot.disengage()
                except Exception:
                    pass
        except Exception:
            pass
        try:
            conn.close()
        except Exception:
            pass

    return status


def _sleep_with_progress(duration_s):
    # 小循环以便能快速响应（比如可拓展为监听中断）
    end = time.time() + duration_s
    while True:
        now = time.time()
        if now >= end:
            break
        time.sleep(min(0.5, end - now))


def _ensure_throttle_applied(conn, sc, desired_throttle=0.7, timeout=10.0, verbose=False):
    """尝试确保对当前活动载具的油门设置被游戏接受。

    优先使用 kRPC streams 订阅 control.throttle、available_thrust、thrust 的回显。
    在分级后调用，函数会在超时前反复写入 throttle 并检查回显或推力变化。
    返回：布尔，表示油门是否被接受。
    """
    try:
        start = time.time()
        vessel = sc.active_vessel
        # 尝试建立 streams
        streams = {}
        try:
            streams['throttle'] = conn.add_stream(getattr, vessel.control, 'throttle')
            streams['available_thrust'] = conn.add_stream(getattr, vessel, 'available_thrust')
            streams['thrust'] = conn.add_stream(getattr, vessel, 'thrust')
        except Exception:
            streams = {}

        while time.time() - start < timeout:
            try:
                vessel = sc.active_vessel
                vessel.control.throttle = desired_throttle
            except Exception:
                pass

            # 读取 stream 回显
            try:
                tval = streams.get('throttle').read() if 'throttle' in streams else getattr(vessel.control, 'throttle', None)
            except Exception:
                try:
                    tval = getattr(vessel.control, 'throttle', None)
                except Exception:
                    tval = None

            try:
                thrust_val = streams.get('thrust').read() if 'thrust' in streams else getattr(vessel, 'thrust', None)
            except Exception:
                thrust_val = None

            try:
                avail = streams.get('available_thrust').read() if 'available_thrust' in streams else getattr(vessel, 'available_thrust', None)
            except Exception:
                avail = None

            # 额外调试输出，帮助定位为何油门被限制
            if verbose:
                try:
                    vs = getattr(vessel, 'name', None)
                except Exception:
                    vs = None
                print(f"_ensure_throttle_applied: vessel={vs} desired={desired_throttle:.3f} tval={tval} thrust={thrust_val} avail={avail}")

            if tval is not None and tval >= desired_throttle * 0.95:
                # 回显已接近目标
                _cleanup_streams(streams)
                return True
            if thrust_val is not None and thrust_val > 0:
                _cleanup_streams(streams)
                return True
            if avail is not None and avail > 0 and tval is not None and tval > 0:
                _cleanup_streams(streams)
                return True

            time.sleep(0.1)

        # 超时仍失败，清理并返回 False
        _cleanup_streams(streams)
        return False
    except Exception:
        return False


def _cleanup_streams(streams):
    try:
        for s in list(streams.values()):
            try:
                s.remove()
            except Exception:
                try:
                    s.close()
                except Exception:
                    pass
    except Exception:
        pass


def _diagnose_engines(vessel):
    """收集并返回当前载具的一些发动机/推进相关诊断信息（尽量安全）。"""
    info = {}
    try:
        info['control_throttle'] = getattr(vessel.control, 'throttle', 'N/A')
    except Exception:
        info['control_throttle'] = 'N/A'
    try:
        info['available_thrust'] = getattr(vessel, 'available_thrust', 'N/A')
    except Exception:
        info['available_thrust'] = 'N/A'
    try:
        info['thrust'] = getattr(vessel, 'thrust', 'N/A')
    except Exception:
        info['thrust'] = 'N/A'
    # 尝试获取部件计数和发动机计数
    try:
        parts = vessel.parts
        info['parts_count'] = getattr(parts, 'count', 'N/A')
        engines = getattr(parts, 'engines', None)
        if engines is not None:
            info['engines_count'] = getattr(engines, 'count', 'N/A')
            # 尝试收集每个 engine 的状态
            engines_list = []
            try:
                for i, eng in enumerate(engines):
                    ed = {'index': i}
                    try:
                        ed['part_name'] = getattr(eng, 'part', None)
                    except Exception:
                        ed['part_name'] = 'N/A'
                    # 常见 engine 属性，安全读取
                    for attr in ('active', 'ignited', 'has_fuel', 'thrust_limit', 'thrust', 'available_thrust', 'fuel_flow', 'flameout'):
                        try:
                            ed[attr] = getattr(eng, attr)
                        except Exception:
                            try:
                                # 有时属性在 engine 上不可用，但可在 part 或 module 上
                                ed[attr] = getattr(eng, attr, 'N/A')
                            except Exception:
                                ed[attr] = 'N/A'
                    engines_list.append(ed)
            except Exception:
                engines_list = 'N/A'
            info['engines'] = engines_list
        else:
            info['engines_count'] = 'N/A'
            info['engines'] = 'N/A'
    except Exception:
        info['parts_count'] = 'N/A'
        info['engines_count'] = 'N/A'
        info['engines'] = 'N/A'

    return info


def _force_apply_throttle(conn, sc, desired_throttle, timeout=5.0, verbose=False):
    """在 _ensure_throttle_applied 失败后尝试更激进地应用油门：
    - 设置每个发动机的 thrust_limit 到目标（如果属性存在）
    - 多次写入 vessel.control.throttle
    返回是否成功检测到油门生效。
    """
    try:
        start = time.time()
        vessel = sc.active_vessel
        # 尝试设置各发动机的 thrust_limit
        try:
            parts = vessel.parts
            engines = getattr(parts, 'engines', None)
            if engines is not None:
                for eng in engines:
                    try:
                        # 有些 engine 对象暴露 thrust_limit 属性
                        if hasattr(eng, 'thrust_limit'):
                            eng.thrust_limit = desired_throttle
                    except Exception:
                        # 忽略单个发动机失败
                        pass
        except Exception:
            pass

        # 重复写入并检测
        while time.time() - start < timeout:
            try:
                vessel = sc.active_vessel
                vessel.control.throttle = float(desired_throttle)
            except Exception:
                pass

            try:
                tval = getattr(vessel.control, 'throttle', None)
            except Exception:
                tval = None
            try:
                thrust_val = getattr(vessel, 'thrust', None)
            except Exception:
                thrust_val = None

            if verbose:
                print(f"_force_apply_throttle: desired={desired_throttle:.3f} tval={tval} thrust={thrust_val}")

            if tval is not None and tval >= desired_throttle * 0.9:
                return True
            if thrust_val is not None and thrust_val > 0:
                return True

            time.sleep(0.1)

        return False
    except Exception:
        return False


if __name__ == "__main__":
    run_flight_sequence(verbose=True)
