from flask import Blueprint, jsonify, request

from hiveos.api import build_sim_agent_for_capability
from hiveos.api.context import authorize_core_action, emit_action_applied, emit_action_failed
from hiveos.config import (
    ALLOW_WRAPPER_PY,
    REQUIRE_WRAPPER_SPEC_SIGNATURE,
    WRAPPER_SPEC_ALLOWLIST,
    WRAPPER_SPEC_SIGNING_KEY,
)
from hiveos.utilities.wrapper_serialization import (
    hydrate_wrapper,
    validate_py_wrapper,
    validate_wrapper_spec_policy,
)


agents_bp = Blueprint("agents", __name__)


@agents_bp.route("/inject_agent", methods=["POST"])
def inject_agent():
    action = "agents.inject"
    try:
        core = authorize_core_action(action, resource_type="session")
    except PermissionError as e:
        return jsonify({"error": str(e)}), 403

    try:
        data = request.json
        agent_source = data.get("agent_source")
        capabilities = data.get("capabilities", [])
        wrapper_spec = data.get("wrapper_spec")
        wrapper_py = data.get("wrapper_py")
        ble_address = data.get("ble_address")

        if wrapper_spec:
            ok, reason = validate_wrapper_spec_policy(
                wrapper_spec,
                allowlist=WRAPPER_SPEC_ALLOWLIST,
                require_signature=REQUIRE_WRAPPER_SPEC_SIGNATURE,
                signing_key=WRAPPER_SPEC_SIGNING_KEY,
            )
            if not ok:
                emit_action_failed(action, resource_type="session", reason=f"wrapper_spec_rejected:{reason}")
                return jsonify({"error": f"wrapper_spec rejected: {reason}"}), 403
            agent = hydrate_wrapper(wrapper_spec)
        elif wrapper_py:
            if not ALLOW_WRAPPER_PY:
                emit_action_failed(action, resource_type="session", reason="wrapper_py_disabled")
                return jsonify(
                    {
                        "error": "wrapper_py ingestion is disabled. Use wrapper_spec or set HIVE_ALLOW_WRAPPER_PY=true."
                    }
                ), 403
            success, result = validate_py_wrapper(wrapper_py)
            if not success:
                emit_action_failed(action, resource_type="session", reason=f"invalid_wrapper:{result}")
                return jsonify({"error": f"Invalid wrapper: {result}"}), 400
            ok, reason = validate_wrapper_spec_policy(
                result,
                allowlist=WRAPPER_SPEC_ALLOWLIST,
                require_signature=REQUIRE_WRAPPER_SPEC_SIGNATURE,
                signing_key=WRAPPER_SPEC_SIGNING_KEY,
            )
            if not ok:
                emit_action_failed(
                    action, resource_type="session", reason=f"wrapper_py_policy_rejected:{reason}"
                )
                return jsonify({"error": f"wrapper_py rejected by wrapper policy: {reason}"}), 403
            agent = hydrate_wrapper(result)
        elif agent_source and capabilities:
            cap = capabilities[0]
            if agent_source == "sim":
                agent, wrapper_spec = build_sim_agent_for_capability(cap, capabilities)
            elif agent_source == "hardware":
                if cap == "display":
                    from hiveos.agents_hardware.oled_agent import OledDisplayAgent

                    if not ble_address:
                        emit_action_failed(action, resource_type="session", reason="missing_ble_address_display")
                        return jsonify({"error": "Missing 'ble_address' for hardware display agent"}), 400
                    agent = OledDisplayAgent(ble_address=ble_address)
                elif cap == "sound":
                    from hiveos.agents_hardware.piezo_agent import PiezoAgent

                    if not ble_address:
                        emit_action_failed(action, resource_type="session", reason="missing_ble_address_sound")
                        return jsonify({"error": "Missing 'ble_address' for hardware sound agent"}), 400
                    agent = PiezoAgent(ble_address=ble_address)
                elif cap == "terminal":
                    from hiveos.agents_hardware.terminal_agent import TerminalAgent

                    if not ble_address:
                        emit_action_failed(action, resource_type="session", reason="missing_ble_address_terminal")
                        return jsonify({"error": "Missing 'ble_address' for hardware terminal agent"}), 400
                    agent = TerminalAgent(ble_address=ble_address)
                else:
                    emit_action_failed(
                        action,
                        resource_type="session",
                        reason=f"unsupported_hardware_capability:{cap}",
                    )
                    return jsonify({"error": f"Unsupported hardware capability: {cap}"}), 400
            else:
                emit_action_failed(action, resource_type="session", reason=f"invalid_agent_source:{agent_source}")
                return jsonify({"error": f"Invalid agent_source: {agent_source}"}), 400
        else:
            emit_action_failed(action, resource_type="session", reason="missing_agent_definition_or_metadata")
            return jsonify({"error": "Missing agent definition or metadata"}), 400

        print("[Test] Wrapper Spec:", wrapper_spec)
        print("[Test] Hydrated Agent:", agent)
        print("[Test] Agent ID:", agent.agent_id)
        print("[Test] Capabilities:", agent.capabilities)

        core.inject_agent(agent)
        emit_action_applied(
            action,
            resource_type="agent",
            resource_id=agent.agent_id,
            payload={"capabilities": list(getattr(agent, "capabilities", []) or [])},
        )
        return jsonify({"message": "Agent injected successfully"}), 200
    except Exception as e:
        emit_action_failed(action, resource_type="session", reason=e)
        return jsonify({"error": f"Failed to inject agent: {e}"}), 500


@agents_bp.route("/inject_agents_for_pending", methods=["POST"])
def inject_agents_for_pending():
    action = "agents.inject_pending"
    try:
        core = authorize_core_action(action, resource_type="session")
    except PermissionError as e:
        return jsonify({"error": str(e)}), 403

    try:
        data = request.get_json(silent=True) or {}
        max_total = max(1, min(int(data.get("max_total", 24)), 100))
        per_capability = max(1, min(int(data.get("per_capability", 2)), 20))
        session_id = request.headers.get("X-Session-ID")

        pending_chunks = [
            c for c in core.chunk_objects.values() if c.session_id == session_id and c.status == "pending"
        ]
        if not pending_chunks:
            emit_action_applied(action, resource_type="session", payload={"injected_total": 0})
            return (
                jsonify(
                    {
                        "status": "ok",
                        "message": "No pending chunks found.",
                        "injected_total": 0,
                        "injected_by_capability": {},
                    }
                ),
                200,
            )

        pending_by_capability = {}
        for chunk in pending_chunks:
            cap = str(chunk.required_capability or "generic")
            pending_by_capability[cap] = pending_by_capability.get(cap, 0) + 1

        active_by_capability = {}
        for runtime_agent in core._runtime_agents.values():
            if runtime_agent.session_id != session_id:
                continue
            caps = list(getattr(runtime_agent.obj, "capabilities", []) or [])
            for cap in caps:
                key = str(cap)
                active_by_capability[key] = active_by_capability.get(key, 0) + 1

        injected_total = 0
        injected_by_capability = {}
        for cap, pending_count in pending_by_capability.items():
            if injected_total >= max_total:
                break
            current = active_by_capability.get(cap, 0)
            desired = min(per_capability, pending_count)
            to_inject = max(desired - current, 0)
            to_inject = min(to_inject, max_total - injected_total)
            if to_inject <= 0:
                continue

            for _ in range(to_inject):
                agent, _ = build_sim_agent_for_capability(cap, [cap])
                core.inject_agent(agent)
            injected_total += to_inject
            injected_by_capability[cap] = to_inject

        emit_action_applied(
            action,
            resource_type="session",
            payload={"injected_total": injected_total, "injected_by_capability": injected_by_capability},
        )
        return (
            jsonify(
                {
                    "status": "ok",
                    "pending_by_capability": pending_by_capability,
                    "active_by_capability": active_by_capability,
                    "injected_by_capability": injected_by_capability,
                    "injected_total": injected_total,
                }
            ),
            200,
        )
    except Exception as e:
        emit_action_failed(action, resource_type="session", reason=e)
        return jsonify({"error": f"Failed to inject pending capability agents: {e}"}), 500


@agents_bp.route("/remove_agent", methods=["POST"])
def remove_agent():
    action = "agents.remove"
    agent_id = None
    try:
        data = request.json
        agent_id = (data or {}).get("agent_id")
        if not agent_id:
            return jsonify({"error": "Missing 'agent_id' field"}), 400
        core = authorize_core_action(action, resource_type="agent", resource_id=agent_id)

        result = core.remove_agent(agent_id)
        if result:
            emit_action_applied(action, resource_type="agent", resource_id=agent_id)
            return jsonify({"message": f"Agent {agent_id} removed"}), 200
        emit_action_failed(action, resource_type="agent", resource_id=agent_id, reason="agent_not_found")
        return jsonify({"error": f"No agent with ID {agent_id} found"}), 404
    except PermissionError as e:
        return jsonify({"error": str(e)}), 403
    except Exception as e:
        import traceback

        traceback.print_exc()
        emit_action_failed(action, resource_type="agent", resource_id=agent_id, reason=e)
        return jsonify({"error": str(e)}), 500


@agents_bp.route("/agent_input", methods=["POST"])
def agent_input():
    action = "agents.input"
    agent_id = None
    try:
        data = request.json
        agent_id = (data or {}).get("agent_id")
        user_input = (data or {}).get("input")

        if not agent_id or user_input is None:
            return jsonify({"error": "Missing 'agent_id' or 'input'"}), 400
        core = authorize_core_action(action, resource_type="agent", resource_id=agent_id)

        agent = core._runtime_agents.get(agent_id).obj if agent_id in core._runtime_agents else None
        if not agent or not hasattr(agent, "receive_input"):
            emit_action_failed(action, resource_type="agent", resource_id=agent_id, reason="agent_not_input_capable")
            return jsonify({"error": "Agent not found or not input-capable"}), 404

        agent.receive_input(user_input)
        emit_action_applied(action, resource_type="agent", resource_id=agent_id)
        return jsonify({"status": "input received"}), 200
    except PermissionError as e:
        return jsonify({"error": str(e)}), 403
    except Exception as e:
        import traceback

        traceback.print_exc()
        emit_action_failed(action, resource_type="agent", resource_id=agent_id, reason=e)
        return jsonify({"error": str(e)}), 500
