import asyncio
import contextlib
import json
from collections.abc import AsyncGenerator
from pathlib import Path

from ntrp.constants import NTRP_TMP_BASE, OFFLOAD_PREVIEW_LINES, OFFLOAD_THRESHOLD
from ntrp.core.async_queue import AsyncQueue
from ntrp.core.models import PendingToolCall, ToolExecutionResult
from ntrp.events.internal import ToolExecuted
from ntrp.events.sse import SSEEvent, ToolCallEvent, ToolResultEvent
from ntrp.logging import get_logger
from ntrp.tools.core.base import ToolResult
from ntrp.tools.core.context import ToolContext, ToolExecution
from ntrp.tools.executor import ToolExecutor
from ntrp.utils import ms_now

_logger = get_logger(__name__)

OFFLOAD_BASE = Path(NTRP_TMP_BASE)


class ToolRunner:
    def __init__(
        self,
        executor: ToolExecutor,
        ctx: ToolContext,
        depth: int,
        parent_id: str | None,
    ):
        self.executor = executor
        self.ctx = ctx
        self.depth = depth
        self.parent_id = parent_id or "root"
        self._offload_counter = 0

    def _maybe_offload(self, tool_name: str, result: ToolResult) -> ToolResult:
        content = result.content
        if len(content) <= OFFLOAD_THRESHOLD:
            return result

        self._offload_counter += 1
        offload_dir = OFFLOAD_BASE / self.ctx.session_id / "results"
        offload_dir.mkdir(parents=True, exist_ok=True)
        offload_path = offload_dir / f"{tool_name}_{self._offload_counter}.txt"
        offload_path.write_text(content, encoding="utf-8")

        compact = self._build_offload_summary(content, offload_path)
        return ToolResult(content=compact, preview=result.preview, data=result.data)

    @staticmethod
    def _build_offload_summary(content: str, offload_path: Path) -> str:
        lines = content.split("\n")
        total = len(lines)
        preview = "\n".join(lines[:OFFLOAD_PREVIEW_LINES])

        return (
            f"{preview}\n\n"
            f"[{total} lines total, {total - OFFLOAD_PREVIEW_LINES} more offloaded → {offload_path}]\n"
            f"Use bash(grep -n 'pattern' '{offload_path}') to find specific content, "
            f"or read_file(path='{offload_path}', offset=N, limit=M) to read a specific section."
        )

    def _display_name(self, name: str) -> str:
        tool = self.executor.registry.get(name)
        return tool.display_name if tool else name

    def _make_start_event(self, call: PendingToolCall) -> ToolCallEvent:
        return ToolCallEvent(
            tool_id=call.tool_call.id,
            name=call.name,
            args=call.args,
            depth=self.depth,
            parent_id=self.parent_id,
            display_name=self._display_name(call.name),
        )

    async def _make_result_event(
        self,
        call: PendingToolCall,
        result: ToolResult,
        duration_ms: int,
    ) -> ToolResultEvent:
        self.ctx.channel.publish(
            ToolExecuted(
                name=call.name,
                duration_ms=duration_ms,
                depth=self.depth,
                is_error=result.is_error,
                run_id=self.ctx.run.run_id,
            )
        )

        tool = self.executor.registry.get(call.name)
        if self.ctx.ledger and not result.is_error and tool and not tool.mutates and not tool.volatile:
            identity = f"{call.name}:{json.dumps(call.args, sort_keys=True)}"
            already_read = await self.ctx.ledger.mark_read(identity)
            if already_read:
                result = ToolResult(
                    content=f"[Already read by another agent in this run]\n{result.content}",
                    preview=result.preview,
                    data=result.data,
                )

        return ToolResultEvent(
            tool_id=call.tool_call.id,
            name=call.name,
            result=result.content,
            preview=result.preview,
            depth=self.depth,
            parent_id=self.parent_id,
            duration_ms=duration_ms,
            data=result.data,
            display_name=self._display_name(call.name),
        )

    def _needs_approval(self, call: PendingToolCall) -> bool:
        tool = self.executor.registry.get(call.name)
        if not tool or not tool.mutates:
            return False
        if self.ctx.skip_approvals:
            return False
        if call.name in self.ctx.auto_approve:
            return False
        return True

    async def _execute_single(self, call: PendingToolCall) -> AsyncGenerator[SSEEvent]:
        yield self._make_start_event(call)

        start_ms = ms_now()
        try:
            execution = ToolExecution(call.tool_call.id, call.name, self.ctx)
            result = await self.executor.execute(call.name, call.args, execution)
            result = self._maybe_offload(call.name, result)
            duration_ms = ms_now() - start_ms
            yield await self._make_result_event(call=call, result=result, duration_ms=duration_ms)
        except Exception as e:
            duration_ms = ms_now() - start_ms
            yield await self._make_result_event(
                call=call,
                result=ToolResult(
                    content=f"Error: {type(e).__name__}: {e}", preview=f"Failed: {type(e).__name__}", is_error=True
                ),
                duration_ms=duration_ms,
            )

    async def _execute_concurrent(self, calls: list[PendingToolCall]) -> AsyncGenerator[SSEEvent]:
        results_queue: AsyncQueue[ToolExecutionResult] = AsyncQueue()

        for call in calls:
            yield self._make_start_event(call)

        async def execute_tool(call: PendingToolCall) -> None:
            start_ms = ms_now()
            try:
                execution = ToolExecution(call.tool_call.id, call.name, self.ctx)
                result = await self.executor.execute(call.name, call.args, execution)
                result = self._maybe_offload(call.name, result)
                duration_ms = ms_now() - start_ms
                results_queue.enqueue(
                    ToolExecutionResult(
                        call=call,
                        content=result.content,
                        preview=result.preview,
                        duration_ms=duration_ms,
                        is_error=result.is_error,
                        data=result.data,
                    )
                )
            except Exception as e:
                duration_ms = ms_now() - start_ms
                results_queue.enqueue(
                    ToolExecutionResult(
                        call=call,
                        content=f"Error: {type(e).__name__}: {e}",
                        preview=f"Failed: {type(e).__name__}",
                        duration_ms=duration_ms,
                        is_error=True,
                    )
                )

        async def run_all() -> None:
            try:
                async with asyncio.TaskGroup() as tg:
                    for call in calls:
                        tg.create_task(execute_tool(call))
            except ExceptionGroup as eg:
                _logger.warning("Tool execution errors: %s", [str(e) for e in eg.exceptions])
            finally:
                results_queue.finish()

        task = asyncio.create_task(run_all())

        try:
            async for r in results_queue:
                yield await self._make_result_event(
                    call=r.call,
                    result=ToolResult(content=r.content, preview=r.preview, is_error=r.is_error, data=r.data),
                    duration_ms=r.duration_ms,
                )
        finally:
            if not task.done():
                task.cancel()
                with contextlib.suppress(asyncio.CancelledError):
                    await task

    async def _execute_sequential(self, calls: list[PendingToolCall]) -> AsyncGenerator[SSEEvent]:
        for call in calls:
            async for event in self._execute_single(call):
                yield event

    async def execute_all(self, calls: list[PendingToolCall]) -> AsyncGenerator[SSEEvent]:
        def partition(pred, items):
            yes, no = [], []
            for item in items:
                (yes if pred(item) else no).append(item)
            return yes, no

        needs_approval, auto_approved = partition(self._needs_approval, calls)

        for batch, executor in [
            (auto_approved, self._execute_concurrent),
            (needs_approval, self._execute_sequential),
        ]:
            if not batch:
                continue
            async for event in executor(batch):
                yield event
