from hiveos.intelligence.tracing import build_trace_event, emit_trace
from hiveos.runtime.states import TaskStatus, ZoneStallState


class CoreToolService:
    def __init__(self, core):
        self.core = core

    def register_builtin_tools(self):
        registry = self.core.tool_registry
        registry.register_tool(
            tool_id="core.get_status_snapshot",
            capability="ai",
            handler=self.get_status_snapshot,
            request_schema={"required_args": []},
            limits={"max_calls_per_minute": 120, "max_cost_per_hour": 1000},
        )
        registry.register_tool(
            tool_id="memory.put",
            capability="ai",
            handler=self.memory_put,
            request_schema={"required_args": ["scope", "key", "value"]},
            limits={"max_calls_per_minute": 240, "max_cost_per_hour": 2000},
        )
        registry.register_tool(
            tool_id="memory.get",
            capability="ai",
            handler=self.memory_get,
            request_schema={"required_args": ["scope", "key"]},
            limits={"max_calls_per_minute": 240, "max_cost_per_hour": 2000},
        )
        registry.register_tool(
            tool_id="memory.query",
            capability="ai",
            handler=self.memory_query,
            request_schema={"required_args": []},
            limits={"max_calls_per_minute": 120, "max_cost_per_hour": 1500},
        )
        registry.register_tool(
            tool_id="sandbox.start_run",
            capability="ai",
            handler=self.sandbox_start_run,
            request_schema={"required_args": ["objective"]},
            limits={"max_calls_per_minute": 30, "max_cost_per_hour": 800},
        )
        registry.register_tool(
            tool_id="sandbox.get_result",
            capability="ai",
            handler=self.sandbox_get_result,
            request_schema={"required_args": ["sandbox_run_id"]},
            limits={"max_calls_per_minute": 120, "max_cost_per_hour": 800},
        )
        registry.register_tool(
            tool_id="sandbox.cancel",
            capability="ai",
            handler=self.sandbox_cancel,
            request_schema={"required_args": ["sandbox_run_id"]},
            limits={"max_calls_per_minute": 30, "max_cost_per_hour": 800},
        )
        registry.register_tool(
            tool_id="terminal.session",
            capability="terminal.stream",
            handler=self.terminal_session,
            request_schema={"required_args": ["target_id", "target_type", "transport", "action"]},
            limits={"max_calls_per_minute": 120, "max_cost_per_hour": 1000},
        )
        registry.register_tool(
            tool_id="terminal.exec",
            capability="terminal.exec",
            handler=self.terminal_exec,
            request_schema={"required_args": ["target_id", "command"]},
            limits={"max_calls_per_minute": 120, "max_cost_per_hour": 2000},
        )
        registry.register_tool(
            tool_id="terminal.transform",
            capability="terminal.stream",
            handler=self.terminal_transform,
            request_schema={"required_args": ["target_id", "module"]},
            limits={"max_calls_per_minute": 120, "max_cost_per_hour": 1000},
        )

    def get_status_snapshot(self, _args, _context):
        blocked_zones = sum(
            1
            for zone in self.core.zones.values()
            if zone.stall_state in {ZoneStallState.BLOCKED, ZoneStallState.FAILED}
        )
        blocked_tasks = sum(
            1 for task in self.core.task_objects.values() if task.status == TaskStatus.BLOCKED
        )
        return {
            "tasks": len(self.core.task_objects),
            "chunks": len(self.core.chunk_objects),
            "zones": len(self.core.zones),
            "agents": len(self.core._runtime_agents),
            "ticks": self.core.tick_counter["count"],
            "session_id": self.core.session_id,
            "runtime_perf": self.core.get_runtime_perf(),
            "blocked_zones": blocked_zones,
            "blocked_tasks": blocked_tasks,
        }

    def memory_put(self, args, context):
        chunk = context.get("chunk")
        agent = context.get("agent")
        run_id = context.get("run_id")

        scope_ref = dict(args.get("scope_ref") or {})
        if chunk is not None:
            scope_ref.setdefault("task_id", chunk.task_id)
            scope_ref.setdefault("chunk_id", chunk.chunk_id)
        if agent is not None:
            scope_ref.setdefault("agent_id", agent.agent_id)

        entry = self.core.memory_store.put(
            {
                "scope": args.get("scope"),
                "scope_ref": scope_ref,
                "key": args.get("key"),
                "value": args.get("value"),
                "content_type": args.get("content_type", "json"),
                "ttl_seconds": args.get("ttl_seconds", 0),
                "tags": args.get("tags", []),
                "source": args.get("source", "ai"),
                "inspect": args.get("inspect", {"visible": True, "redaction": "none", "notes": ""}),
            }
        )
        emit_trace(
            build_trace_event(
                event_type="memory_put",
                run_id=run_id or f"{entry['entry_id']}:memory_put",
                outcome="success",
                payload={"entry_id": entry["entry_id"], "scope": entry["scope"], "key": entry["key"]},
                agent_id=(agent.agent_id if agent else None),
                task_id=(chunk.task_id if chunk else None),
                chunk_id=(chunk.chunk_id if chunk else None),
                zone_id=(agent.zone if agent else None),
            )
        )
        return {
            "entry_id": entry["entry_id"],
            "scope": entry["scope"],
            "key": entry["key"],
            "expires_at": entry["expires_at"],
        }

    def memory_get(self, args, context):
        chunk = context.get("chunk")
        agent = context.get("agent")
        run_id = context.get("run_id")
        scope_ref = dict(args.get("scope_ref") or {})
        if chunk is not None:
            scope_ref.setdefault("task_id", chunk.task_id)
            scope_ref.setdefault("chunk_id", chunk.chunk_id)
        if agent is not None:
            scope_ref.setdefault("agent_id", agent.agent_id)

        entry = self.core.memory_store.get(
            scope=args.get("scope"),
            key=args.get("key"),
            scope_ref=scope_ref,
            include_expired=bool(args.get("include_expired", False)),
        )
        emit_trace(
            build_trace_event(
                event_type="memory_get",
                run_id=run_id or f"{args.get('key', 'key')}:memory_get",
                outcome="success",
                payload={
                    "scope": args.get("scope"),
                    "key": args.get("key"),
                    "hit": bool(entry),
                    "entry_id": (entry or {}).get("entry_id"),
                },
                agent_id=(agent.agent_id if agent else None),
                task_id=(chunk.task_id if chunk else None),
                chunk_id=(chunk.chunk_id if chunk else None),
                zone_id=(agent.zone if agent else None),
            )
        )
        return entry

    def memory_query(self, args, context):
        chunk = context.get("chunk")
        agent = context.get("agent")
        run_id = context.get("run_id")
        filters = dict(args.get("filters") or {})
        scope_ref = dict(filters.get("scope_ref") or {})
        if chunk is not None:
            scope_ref.setdefault("task_id", chunk.task_id)
            scope_ref.setdefault("chunk_id", chunk.chunk_id)
        if agent is not None:
            scope_ref.setdefault("agent_id", agent.agent_id)
        filters["scope_ref"] = scope_ref

        entries = self.core.memory_store.query(
            scope=args.get("scope"),
            filters=filters,
            include_expired=bool(args.get("include_expired", False)),
        )
        emit_trace(
            build_trace_event(
                event_type="memory_query",
                run_id=run_id or "memory_query",
                outcome="success",
                payload={"scope": args.get("scope"), "result_count": len(entries)},
                agent_id=(agent.agent_id if agent else None),
                task_id=(chunk.task_id if chunk else None),
                chunk_id=(chunk.chunk_id if chunk else None),
                zone_id=(agent.zone if agent else None),
            )
        )
        return entries

    def sandbox_start_run(self, args, context):
        chunk = context.get("chunk")
        agent = context.get("agent")
        run_id = context.get("run_id")

        spec = {
            "objective": args.get("objective"),
            "constraints": args.get("constraints", {}),
            "snapshot_ref": args.get("snapshot_ref"),
            "task_id": (chunk.task_id if chunk else args.get("task_id")),
            "chunk_id": (chunk.chunk_id if chunk else args.get("chunk_id")),
            "agent_id": (agent.agent_id if agent else args.get("agent_id")),
        }
        try:
            sandbox_run_id = self.core.sandbox_runtime.start_run(spec)
        except Exception as exc:
            emit_trace(
                build_trace_event(
                    event_type="sandbox_run_rejected",
                    run_id=run_id or "sandbox",
                    outcome="denied",
                    payload={"reason": str(exc), "objective": spec.get("objective")},
                    agent_id=(agent.agent_id if agent else None),
                    task_id=(chunk.task_id if chunk else None),
                    chunk_id=(chunk.chunk_id if chunk else None),
                    zone_id=(agent.zone if agent else None),
                )
            )
            return {"status": "rejected", "reason": str(exc)}
        emit_trace(
            build_trace_event(
                event_type="sandbox_run_started",
                run_id=run_id or sandbox_run_id,
                outcome="success",
                payload={"sandbox_run_id": sandbox_run_id, "objective": spec.get("objective")},
                agent_id=(agent.agent_id if agent else None),
                task_id=(chunk.task_id if chunk else None),
                chunk_id=(chunk.chunk_id if chunk else None),
                zone_id=(agent.zone if agent else None),
            )
        )
        return {"sandbox_run_id": sandbox_run_id, "state": "queued"}

    def sandbox_get_result(self, args, context):
        chunk = context.get("chunk")
        agent = context.get("agent")
        run_id = context.get("run_id")
        sandbox_run_id = args.get("sandbox_run_id")
        result = self.core.sandbox_runtime.get_result(sandbox_run_id)
        emit_trace(
            build_trace_event(
                event_type="sandbox_run_observed",
                run_id=run_id or sandbox_run_id,
                outcome="success",
                payload={
                    "sandbox_run_id": sandbox_run_id,
                    "found": bool(result),
                    "state": (result or {}).get("state"),
                },
                agent_id=(agent.agent_id if agent else None),
                task_id=(chunk.task_id if chunk else None),
                chunk_id=(chunk.chunk_id if chunk else None),
                zone_id=(agent.zone if agent else None),
            )
        )
        return result

    def sandbox_cancel(self, args, context):
        chunk = context.get("chunk")
        agent = context.get("agent")
        run_id = context.get("run_id")
        sandbox_run_id = args.get("sandbox_run_id")
        canceled = self.core.sandbox_runtime.cancel(sandbox_run_id)
        outcome = "success" if canceled else "denied"
        emit_trace(
            build_trace_event(
                event_type="sandbox_run_canceled",
                run_id=run_id or sandbox_run_id,
                outcome=outcome,
                payload={"sandbox_run_id": sandbox_run_id, "canceled": canceled},
                agent_id=(agent.agent_id if agent else None),
                task_id=(chunk.task_id if chunk else None),
                chunk_id=(chunk.chunk_id if chunk else None),
                zone_id=(agent.zone if agent else None),
            )
        )
        return {"sandbox_run_id": sandbox_run_id, "canceled": canceled}

    def terminal_session(self, args, context):
        action = str(args.get("action") or "").strip().lower()
        if action == "open":
            return self.core.terminal_service.open_target(
                {
                    "target_id": args.get("target_id"),
                    "target_type": args.get("target_type"),
                    "session_id": self.core.session_id,
                    "agent_id": args.get("agent_id"),
                    "transport": args.get("transport"),
                    "meta": args.get("meta"),
                }
            )
        if action == "close":
            closed = self.core.terminal_service.close_target(str(args.get("target_id") or ""))
            return {"target_id": args.get("target_id"), "closed": closed}
        raise ValueError(f"invalid_terminal_session_action:{action}")

    def terminal_exec(self, args, context):
        return self.core.terminal_service.run_command(
            target_id=str(args.get("target_id") or ""),
            command=str(args.get("command") or ""),
            cwd=args.get("cwd"),
            env=args.get("env"),
            timeout_sec=int(args.get("timeout_sec") or 30),
        )

    def terminal_transform(self, args, context):
        return self.core.terminal_service.set_transform(
            target_id=str(args.get("target_id") or ""),
            module=str(args.get("module") or ""),
            config=args.get("config") or {},
        )
