from __future__ import annotations

import asyncio
import json
import logging
import signal
import sys
from dataclasses import dataclass
from typing import Literal, Optional

from clawagents.providers.llm import LLMProvider, LLMMessage, LLMResponse
from clawagents.tools.registry import ToolRegistry, ParsedToolCall

logger = logging.getLogger(__name__)

AgentStatus = Literal["understanding", "acting", "verifying", "done", "error"]


@dataclass
class AgentState:
    messages: list[LLMMessage]
    current_task: str
    status: AgentStatus
    result: str
    iterations: int
    max_iterations: int
    tool_calls: int


BASE_SYSTEM_PROMPT = """You are a ClawAgent, an AI assistant that helps users accomplish tasks using tools. You respond with text and tool calls.

## Core Behavior
- Be concise and direct. Don't over-explain unless asked.
- NEVER add unnecessary preamble ("Sure!", "Great question!", "I'll now...").
- If the request is ambiguous, ask questions before acting.

## Doing Tasks
When the user asks you to do something:
1. **Understand first** — read relevant files, check existing patterns.
2. **Act** — use tools to implement the solution. Work quickly but accurately.
3. **Verify** — check your work against what was asked, not against your own output.

Keep working until the task is fully complete."""


# ─── Helpers ──────────────────────────────────────────────────────────────


def _make_on_chunk():
    """Return a sync callback that streams text to stdout."""
    def on_chunk(chunk: str) -> None:
        sys.stdout.write(chunk)
        sys.stdout.flush()
    return on_chunk


def _get_on_chunk(streaming: bool):
    """Return on_chunk callback if streaming, else None."""
    return _make_on_chunk() if streaming else None


def _log_partial(response: LLMResponse, node: str) -> None:
    if response.partial:
        logger.warning("  [%s] Received partial response (stream was interrupted)", node)
        print(f"  [{node}] ⚠ Received partial response (stream was interrupted)",
              file=sys.stderr)


# ─── Node Functions ───────────────────────────────────────────────────────


async def understand_node(
    state: AgentState,
    llm: LLMProvider,
    tool_desc: str,
    system_prompt: Optional[str] = None,
    cancel_event: asyncio.Event | None = None,
    streaming: bool = True,
) -> AgentState:
    print(f'  [Understand] Analyzing task: "{state.current_task}"')

    prompt_to_use = system_prompt or BASE_SYSTEM_PROMPT
    messages = [
        LLMMessage(role="system", content=f"{prompt_to_use}\n\n{tool_desc}"),
        *state.messages,
        LLMMessage(
            role="user",
            content=(
                f'Analyze this task and create a brief plan of action. '
                f'Task: "{state.current_task}". '
                f'List the tools you\'ll use and in what order. Be concise (2-3 sentences).'
            ),
        ),
    ]

    response = await llm.chat(
        messages, on_chunk=_get_on_chunk(streaming), cancel_event=cancel_event
    )
    if streaming:
        print()
    _log_partial(response, "Understand")

    new_messages = list(state.messages)
    new_messages.append(
        LLMMessage(role="assistant", content=f"[Understanding] {response.content}")
    )

    return AgentState(
        messages=new_messages,
        current_task=state.current_task,
        status="acting",
        result=state.result,
        iterations=state.iterations,
        max_iterations=state.max_iterations,
        tool_calls=state.tool_calls,
    )


MAX_TOOL_ROUNDS_PER_ACT = 5


def _log_tool_call(call: ParsedToolCall) -> None:
    print(f"    -> Tool: {call.tool_name}({json.dumps(call.args)})")


async def act_node(
    state: AgentState,
    llm: LLMProvider,
    tools: ToolRegistry,
    tool_desc: str,
    system_prompt: Optional[str] = None,
    cancel_event: asyncio.Event | None = None,
    streaming: bool = True,
) -> AgentState:
    print("  [Act] Executing planned actions...")

    prompt_to_use = system_prompt or BASE_SYSTEM_PROMPT
    messages = [
        LLMMessage(role="system", content=f"{prompt_to_use}\n\n{tool_desc}"),
        *state.messages,
        LLMMessage(
            role="user",
            content=(
                "Now execute the task using tools. If a tool is needed, respond "
                "with ONLY the JSON tool call block. If no tools are needed, "
                "provide the final result directly."
            ),
        ),
    ]

    current_messages = list(state.messages)
    total_tool_calls = state.tool_calls
    last_result = ""
    on_chunk = _get_on_chunk(streaming)

    for _ in range(MAX_TOOL_ROUNDS_PER_ACT):
        if cancel_event and cancel_event.is_set():
            print("  [Act] Cancelled by user")
            break

        response = await llm.chat(messages, on_chunk=on_chunk, cancel_event=cancel_event)
        if streaming:
            print()
        _log_partial(response, "Act")

        if response.partial and not response.content.strip():
            last_result = "[interrupted — no content received]"
            current_messages.append(
                LLMMessage(role="assistant", content=f"[Action] {last_result}")
            )
            break

        tool_calls = tools.parse_tool_calls(response.content)

        if not tool_calls:
            last_result = response.content
            current_messages.append(
                LLMMessage(role="assistant", content=f"[Action] {response.content}")
            )
            break

        if len(tool_calls) == 1:
            call = tool_calls[0]
            _log_tool_call(call)

            tool_result = await tools.execute_tool(call.tool_name, call.args)
            total_tool_calls += 1

            tool_output = (
                tool_result.output if tool_result.success else f"Error: {tool_result.error}"
            )
            status_tag = "OK" if tool_result.success else "FAIL"
            print(f"    <- {status_tag}: {tool_output[:100]}...")

            current_messages.append(
                LLMMessage(
                    role="assistant",
                    content=f"[Tool Call] {call.tool_name}: {json.dumps(call.args)}",
                )
            )
            current_messages.append(
                LLMMessage(role="user", content=f"[Tool Result] {tool_output}")
            )
        else:
            print(f"    [Parallel] Executing {len(tool_calls)} tool calls concurrently")
            for call in tool_calls:
                _log_tool_call(call)

            results = await tools.execute_tools_parallel(tool_calls)
            total_tool_calls += len(tool_calls)

            call_summaries: list[str] = []
            for j, (call, result) in enumerate(zip(tool_calls, results)):
                output = result.output if result.success else f"Error: {result.error}"
                status_tag = "OK" if result.success else "FAIL"
                print(f"    <- [{j + 1}/{len(tool_calls)}] {status_tag}: {output[:80]}...")
                call_summaries.append(f"{call.tool_name}({json.dumps(call.args)}) => {output}")

            current_messages.append(
                LLMMessage(
                    role="assistant",
                    content=f"[Parallel Tool Calls] {', '.join(c.tool_name for c in tool_calls)}",
                )
            )
            current_messages.append(
                LLMMessage(
                    role="user",
                    content=f"[Tool Results]\n" + "\n".join(call_summaries),
                )
            )

            last_result = "\n".join(call_summaries)

        messages = [
            LLMMessage(role="system", content=f"{prompt_to_use}\n\n{tool_desc}"),
            *current_messages,
            LLMMessage(
                role="user",
                content=(
                    "Continue executing. Use another tool if needed, "
                    "or provide the final result if done."
                ),
            ),
        ]

    return AgentState(
        messages=current_messages,
        current_task=state.current_task,
        status="verifying",
        result=last_result,
        iterations=state.iterations,
        max_iterations=state.max_iterations,
        tool_calls=total_tool_calls,
    )


async def verify_node(
    state: AgentState,
    llm: LLMProvider,
    tool_desc: str,
    system_prompt: Optional[str] = None,
    cancel_event: asyncio.Event | None = None,
    streaming: bool = True,
) -> AgentState:
    print("  [Verify] Checking results...")

    prompt_to_use = system_prompt or BASE_SYSTEM_PROMPT
    messages = [
        LLMMessage(role="system", content=prompt_to_use),
        *state.messages,
        LLMMessage(
            role="user",
            content=(
                'Verify the quality of the result. Is the task complete and correct? '
                'Reply with ONLY "PASS" if the result is satisfactory, '
                'or "RETRY: <reason>" if it needs improvement.'
            ),
        ),
    ]

    response = await llm.chat(
        messages, on_chunk=_get_on_chunk(streaming), cancel_event=cancel_event
    )
    if streaming:
        print()
    _log_partial(response, "Verify")

    verdict = response.content.strip()

    if verdict.startswith("PASS") or state.iterations >= state.max_iterations:
        return AgentState(
            messages=state.messages,
            current_task=state.current_task,
            status="done",
            result=state.result,
            iterations=state.iterations + 1,
            max_iterations=state.max_iterations,
            tool_calls=state.tool_calls,
        )

    print(f"  [Verify] Needs improvement: {verdict}")
    new_messages = list(state.messages)
    new_messages.append(
        LLMMessage(role="assistant", content=f"[Verification] {verdict}")
    )

    return AgentState(
        messages=new_messages,
        current_task=state.current_task,
        status="understanding",
        result=state.result,
        iterations=state.iterations + 1,
        max_iterations=state.max_iterations,
        tool_calls=state.tool_calls,
    )


# ─── Graph Executor ──────────────────────────────────────────────────────


async def run_agent_graph(
    task: str,
    llm: LLMProvider,
    tools: Optional[ToolRegistry] = None,
    system_prompt: Optional[str] = None,
    max_iterations: int = 3,
    streaming: bool = True,
) -> AgentState:
    registry = tools or ToolRegistry()
    tool_desc = registry.describe_for_llm()

    state = AgentState(
        messages=[LLMMessage(role="user", content=task)],
        current_task=task,
        status="understanding",
        result="",
        iterations=0,
        max_iterations=max_iterations,
        tool_calls=0,
    )

    # Cancellation support: abort streams on SIGINT (Ctrl+C)
    cancel_event = asyncio.Event()
    loop = asyncio.get_running_loop()
    prev_handler = None

    def _on_sigint() -> None:
        print("\n  [Agent] Interrupted — aborting current stream...")
        cancel_event.set()

    try:
        prev_handler = loop.add_signal_handler(signal.SIGINT, _on_sigint)  # type: ignore[assignment]
    except (NotImplementedError, OSError):
        pass  # Windows / unsupported loop — fall back to KeyboardInterrupt

    print(f'\n🦞 ClawAgent starting task: "{task}"')
    print(f"   Provider: {llm.name} | Max iterations: {max_iterations} | Streaming: {streaming}")
    tool_names = [t.name for t in registry.list()]
    print(f"   Tools: {', '.join(tool_names) if tool_names else 'none'}\n")

    try:
        while state.status not in ("done", "error"):
            if cancel_event.is_set():
                print("  [Agent] Aborting — user cancelled")
                state.status = "done"
                state.result = state.result or "[cancelled]"
                break

            try:
                if state.status == "understanding":
                    state = await understand_node(
                        state, llm, tool_desc,
                        system_prompt=system_prompt, cancel_event=cancel_event,
                        streaming=streaming,
                    )
                elif state.status == "acting":
                    state = await act_node(
                        state, llm, registry, tool_desc,
                        system_prompt=system_prompt, cancel_event=cancel_event,
                        streaming=streaming,
                    )
                elif state.status == "verifying":
                    state = await verify_node(
                        state, llm, tool_desc,
                        system_prompt=system_prompt, cancel_event=cancel_event,
                        streaming=streaming,
                    )
                else:
                    state.status = "error"
                    state.result = f"Unknown status: {state.status}"
            except KeyboardInterrupt:
                print("\n  [Agent] KeyboardInterrupt — stopping gracefully")
                state.status = "done"
                state.result = state.result or "[interrupted]"
                break
            except asyncio.CancelledError:
                print("\n  [Agent] Task cancelled")
                state.status = "done"
                state.result = state.result or "[cancelled]"
                break
            except Exception as err:
                logger.exception("Error in %s node", state.status)
                print(f"  Error in {state.status} node: {err}")
                state.status = "error"
                state.result = str(err)
    finally:
        try:
            loop.remove_signal_handler(signal.SIGINT)
        except (NotImplementedError, OSError):
            pass

    print(f"\n🦞 ClawAgent finished. Status: {state.status} | Tool calls: {state.tool_calls}")
    return state
