"""AgenticLoop - Executes the Agentic Loop: LLM + Tools iteration."""

import uuid
import logging
from typing import List, Dict, AsyncGenerator

from luxx.tools.executor import ToolExecutor
from luxx.services.llm_client import LLMClient
from luxx.services.stream_context import StreamContext, _sse_event
from luxx.services.llm_response import ParsedDelta, StepType

logger = logging.getLogger(__name__)

MAX_ITERATIONS = 10


class AgenticLoop:
    def __init__(self, tool_executor: ToolExecutor):
        self.tool_executor = tool_executor

    async def execute(
        self,
        llm: LLMClient,
        model: str,
        messages: List[Dict],
        tools: list,
        temperature: float,
        max_tokens: int,
        thinking_enabled: bool,
        context: 'StreamContext',
        tool_context: dict = None
    ) -> AsyncGenerator[str, None]:
        total_usage = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}

        for iteration in range(MAX_ITERATIONS):
            context.reset()
            has_error = False

            async for delta in llm.stream_call(
                model=model,
                messages=messages,
                tools=tools,
                temperature=temperature,
                max_tokens=max_tokens,
                thinking_enabled=thinking_enabled
            ):
                events = self._process_delta(delta, context, total_usage)
                for event in events:
                    yield event

                if not delta.has_content() and not delta.is_complete:
                    has_error = True
                    break

            if has_error:
                break

            # Flush remaining content on complete
            if delta.is_complete:
                for event in self._flush_remaining(context):
                    yield event

            context.finalize_step()

            if context.tool_calls_list:
                for event in self._execute_tools(context, messages, tool_context):
                    yield event
                continue

            for event in self._complete(context, total_usage):
                yield event
            return

        if not has_error:
            yield _sse_event("error", {"content": "Exceeded maximum tool call iterations"})

    def _process_delta(
        self,
        delta: ParsedDelta,
        ctx: 'StreamContext',
        total_usage: dict
    ) -> List[str]:
        events = []

        if delta.usage:
            total_usage.update({
                "prompt_tokens": delta.usage.get("prompt_tokens", 0),
                "completion_tokens": delta.usage.get("completion_tokens", 0),
                "total_tokens": delta.usage.get("total_tokens", 0)
            })

        if delta.content:
            result = ctx.process_content(delta.content)
            if result["should_emit"]:
                # Only emit if there's content
                if result["thinking"]:
                    ctx.full_thinking += result["thinking"]
                    ctx.start_step(StepType.THINKING)
                    events.append(ctx.emit_thinking())

                if result["text"]:
                    ctx.full_content += result["text"]
                    ctx.start_step(StepType.TEXT)
                    events.append(ctx.emit_text())

                # Clear buffers after emit
                ctx._thinking_buf = ""
                ctx._text_buf = ""

        if delta.has_tool_call():
            ctx.accumulate_tool_call(delta.tool_call)

        return events

    def _execute_tools(self, ctx: 'StreamContext', messages: list,
                       tool_context: dict = None) -> List[str]:
        events = []

        for event in ctx.emit_tool_calls():
            events.append(event)

        tool_results = self.tool_executor.process_tool_calls_parallel(
            ctx.tool_calls_list, tool_context or {}
        )

        tool_ids = [tc.get("id") for tc in ctx.tool_calls_list]
        tool_step_ids = [
            s.id for s in ctx.all_steps
            if s.type == StepType.TOOL_CALL and s.id_ref in tool_ids
        ]

        for i, (tr, tc) in enumerate(zip(tool_results, ctx.tool_calls_list)):
            ref_id = tool_step_ids[i] if i < len(tool_step_ids) else f"step-{len(ctx.all_steps) - len(tool_results) + i}"
            _, event = ctx.emit_tool_result(tr, ref_id)
            events.append(event)

        messages.append({
            "role": "assistant",
            "content": ctx.full_content or "",
            "tool_calls": ctx.tool_calls_list
        })
        messages.extend(ctx.all_tool_results[-len(tool_results):])

        return events

    def _flush_remaining(self, ctx: 'StreamContext') -> List[str]:
        """Flush remaining buffers on complete."""
        events = []
        thinking, text = ctx.flush()
        if thinking:
            ctx.full_thinking += thinking
            ctx.start_step(StepType.THINKING)
            events.append(ctx.emit_thinking())
            ctx.finalize_step()
        if text:
            ctx.full_content += text
            ctx.start_step(StepType.TEXT)
            events.append(ctx.emit_text())
            ctx.finalize_step()
        return events

    def _complete(self, ctx: 'StreamContext', total_usage: dict) -> List[str]:
        # Note: buffers already flushed in _flush_remaining or _process_delta
        token_count = total_usage.get("completion_tokens") or len(ctx.full_content) // 4
        msg_id = str(uuid.uuid4())
        logger.info(f"[TOKEN] usage={total_usage}, count={token_count}")

        ctx.set_completion(msg_id, token_count, total_usage)

        return [_sse_event("done", {
            "message_id": msg_id,
            "token_count": token_count,
            "usage": total_usage
        })]