"""Tool-backed background job execution.

This is the "job" execution mode described in RFC_AGENT_BACKGROUND_TASKS:
- A single tool call runs asynchronously in a background task.
- No planner loop is started.

Notes:
- Pause/resume is intentionally not supported in job mode. Tools that require
  pause (OAuth, confirmations) should be executed in "subagent" mode instead.
"""

from __future__ import annotations

import json
from collections import ChainMap
from collections.abc import Callable, Mapping, MutableMapping
from types import MappingProxyType
from typing import Any

from pydantic import BaseModel

from penguiflow.artifacts import ArtifactStore, NoOpArtifactStore, ScopedArtifacts
from penguiflow.catalog import NodeSpec
from penguiflow.planner.context import PlannerPauseReason, ToolContext

from .models import ContextPatch, TaskContextSnapshot, TaskType, UpdateType
from .session import TaskPipeline, TaskResult, TaskRuntime


class ToolJobContext(ToolContext):
    def __init__(
        self,
        *,
        llm_context: Mapping[str, Any],
        tool_context: dict[str, Any],
        artifacts: ArtifactStore | None,
        state_store: object | None = None,
        checkpoint_publisher: Callable[[Mapping[str, Any]], None] | None = None,
    ) -> None:
        self._llm_context = dict(llm_context)
        self._tool_context = dict(tool_context)
        self._artifacts_store = artifacts or NoOpArtifactStore()
        self._state_store = state_store
        self._checkpoint_publisher = checkpoint_publisher
        self._kv = None
        self._scoped_artifacts = ScopedArtifacts(
            self._artifacts_store,
            tenant_id=str(tool_context["tenant_id"]) if tool_context.get("tenant_id") is not None else None,
            user_id=str(tool_context["user_id"]) if tool_context.get("user_id") is not None else None,
            session_id=str(tool_context["session_id"]) if tool_context.get("session_id") is not None else None,
            trace_id=str(tool_context["trace_id"]) if tool_context.get("trace_id") is not None else None,
        )

    @property
    def llm_context(self) -> Mapping[str, Any]:
        return MappingProxyType(self._llm_context)

    @property
    def tool_context(self) -> dict[str, Any]:
        return self._tool_context

    @property
    def meta(self) -> MutableMapping[str, Any]:
        # Deprecated; keep minimal but mutable for ToolContext compatibility.
        return ChainMap(self._tool_context, self._llm_context)

    @property
    def _artifacts(self) -> ArtifactStore:
        return self._artifacts_store

    @property
    def artifacts(self) -> ScopedArtifacts:
        """Scoped artifact facade for tool developers."""
        return self._scoped_artifacts

    @property
    def kv(self):
        from .session_kv import SessionKVFacade

        if self._kv is None:
            self._kv = SessionKVFacade(
                state_store=self._state_store,
                artifacts=self._artifacts,
                tool_context=self._tool_context,
                emit_checkpoint_update=self._checkpoint_publisher,
            )
        return self._kv

    async def pause(
        self,
        reason: PlannerPauseReason,
        payload: Mapping[str, Any] | None = None,
    ) -> Any:
        _ = reason, payload
        raise RuntimeError("pause_not_supported_in_tool_job")

    async def emit_chunk(
        self,
        stream_id: str,
        seq: int,
        text: str,
        *,
        done: bool = False,
        meta: Mapping[str, Any] | None = None,
    ) -> None:
        _ = stream_id, seq, text, done, meta
        return None

    async def emit_artifact(
        self,
        stream_id: str,
        chunk: Any,
        *,
        done: bool = False,
        artifact_type: str | None = None,
        meta: Mapping[str, Any] | None = None,
    ) -> None:
        _ = stream_id, chunk, done, artifact_type, meta
        return None


def _extract_digest(payload: Any) -> list[str]:
    if payload is None:
        return []
    if isinstance(payload, str):
        return [payload[:5000]]
    if isinstance(payload, Mapping):
        for key in ("raw_answer", "answer", "text", "content", "message"):
            if key in payload and payload[key] is not None:
                return [str(payload[key])[:5000]]
    return [str(payload)[:5000]]


def _infer_artifact_type(payload: Any) -> str | None:
    if isinstance(payload, Mapping):
        component = payload.get("component")
        if isinstance(component, str) and component:
            return component
        payload_type = payload.get("type")
        if isinstance(payload_type, str) and payload_type:
            return payload_type
    return None


def _extract_compact_metadata(payload: Any) -> dict[str, Any]:
    if not isinstance(payload, Mapping):
        return {}
    allowed_keys = {
        "type",
        "chart_type",
        "chart_id",
        "analysis_id",
        "title",
        "metric",
        "unit",
    }
    compact: dict[str, Any] = {}
    for key in allowed_keys:
        value = payload.get(key)
        if isinstance(value, (str, int, float, bool)):
            compact[key] = value
    return compact


async def _extract_artifacts_from_observation(
    *,
    node_name: str,
    out_model: type[BaseModel],
    observation: Mapping[str, Any],
    artifacts: ScopedArtifacts,
) -> list[dict[str, Any]]:
    collected: list[dict[str, Any]] = []
    for field_name, field_info in out_model.model_fields.items():
        extra = field_info.json_schema_extra
        if not isinstance(extra, Mapping) or not extra.get("artifact"):
            continue
        if field_name not in observation:
            continue
        value = observation.get(field_name)
        if value is None:
            continue

        async def _store(
            item: Any,
            *,
            item_index: int | None,
            _field_name: str = field_name,
            _node_name: str = node_name,
        ) -> None:
            serialized = json.dumps(item, ensure_ascii=False)
            ref = await artifacts.upload(
                serialized,
                mime_type="application/json",
                filename=f"{_node_name}.{_field_name}.json",
                namespace=f"tool_artifact.{_node_name}.{_field_name}",
                meta={
                    "node": _node_name,
                    "field": _field_name,
                    "item_index": item_index,
                },
            )
            artifact_type = _infer_artifact_type(item)
            compact_meta = _extract_compact_metadata(item)
            stub: dict[str, Any] = {
                "artifact": ref.model_dump(mode="json"),
                **compact_meta,
            }
            if artifact_type:
                stub["type"] = artifact_type
            entry: dict[str, Any] = {
                "node": _node_name,
                "field": _field_name,
                "artifact": stub,
            }
            if item_index is not None:
                entry["item_index"] = item_index
            collected.append(entry)

        if isinstance(value, list):
            for idx, item in enumerate(value):
                if item is None:
                    continue
                await _store(item, item_index=idx)
        else:
            await _store(value, item_index=None)
    return collected


def build_tool_job_pipeline(
    *,
    spec: NodeSpec,
    args_payload: dict[str, Any],
    artifacts: ArtifactStore | None = None,
) -> TaskPipeline:
    """Build a TaskPipeline that runs a single tool call."""

    async def _pipeline(runtime: TaskRuntime) -> TaskResult:
        runtime.emit_update(
            UpdateType.TOOL_CALL,
            {
                "tool_name": spec.name,
                "args": json.loads(json.dumps(args_payload, ensure_ascii=False)),
                "mode": "background_job",
            },
        )
        snapshot: TaskContextSnapshot = runtime.context_snapshot
        # SessionStateStore may be an adapter over a core StateStore; Session KV needs
        # access to the core store for save_memory_state/load_memory_state.
        kv_store = getattr(runtime.session._state_store, "_store", runtime.session._state_store)
        tool_call_id = f"job:{runtime.state.task_id}"
        trace_id = (
            runtime.state.trace_id
            or runtime.context_snapshot.trace_id
            or tool_call_id
        )
        def _publish_checkpoint(payload: Mapping[str, Any]) -> None:
            runtime.emit_update(UpdateType.CHECKPOINT, dict(payload))
            return None

        ctx = ToolJobContext(
            llm_context=snapshot.llm_context or {},
            tool_context={
                **(snapshot.tool_context or {}),
                "session_id": snapshot.session_id,
                "task_id": runtime.state.task_id,
                "trace_id": trace_id,
                "_current_tool_name": spec.name,
                "_current_tool_call_id": tool_call_id,
            },
            artifacts=artifacts,
            state_store=kv_store,
            checkpoint_publisher=_publish_checkpoint,
        )
        parsed_args = spec.args_model.model_validate(args_payload)
        result = await spec.node.func(parsed_args, ctx)
        observation: BaseModel = spec.out_model.model_validate(result)
        payload = observation.model_dump(mode="json")
        extracted_artifacts = []
        if isinstance(payload, Mapping):
            extracted_artifacts = await _extract_artifacts_from_observation(
                node_name=spec.name,
                out_model=spec.out_model,
                observation=payload,
                artifacts=ctx.artifacts,
            )
        patch = ContextPatch(
            task_id=runtime.state.task_id,
            spawned_from_event_id=runtime.context_snapshot.spawned_from_event_id,
            source_context_version=runtime.context_snapshot.context_version,
            source_context_hash=runtime.context_snapshot.context_hash,
            digest=_extract_digest(payload),
            artifacts=extracted_artifacts,
        )
        return TaskResult(
            payload=payload,
            context_patch=patch if runtime.state.task_type == TaskType.BACKGROUND else None,
            digest=patch.digest,
            artifacts=extracted_artifacts,
            sources=[],
            notification=None,
        )

    return _pipeline


__all__ = ["ToolJobContext", "build_tool_job_pipeline"]
