from __future__ import annotations

import asyncio
import copy
import json
from dataclasses import dataclass
from typing import Any, Mapping, Sequence, cast

from .event_stream import EventStream
from .types import (
    AgentContext,
    AgentEvent,
    AgentLoopConfig,
    AgentMessage,
    AgentTool,
    AgentToolResult,
    AssistantMessage,
    Context,
    Message,
    Model,
    StreamFn,
    ToolCall,
    ToolResultMessage,
    clone_message,
    get_value,
    maybe_await,
    message_content,
    message_role,
    model_provider,
    now_ms,
    stop_reason,
)


def stream_simple(model: Model, context: Context, options: Mapping[str, Any] | None = None):
    raise NotImplementedError(
        "stream_simple provider integrations are not part of pidevkit.agent. Pass `stream_fn` explicitly."
    )


def agent_loop(
    prompts: Sequence[AgentMessage],
    context: AgentContext | Mapping[str, Any],
    config: AgentLoopConfig | Mapping[str, Any],
    signal: Any | None = None,
    stream_fn: StreamFn | None = None,
) -> EventStream[AgentEvent, list[AgentMessage]]:
    stream = _create_agent_stream()

    async def runner() -> None:
        new_messages: list[AgentMessage] = list(prompts)
        current_context = _build_context(context, append=prompts)

        stream.push({"type": "agent_start"})
        stream.push({"type": "turn_start"})
        for prompt in prompts:
            stream.push({"type": "message_start", "message": prompt})
            stream.push({"type": "message_end", "message": prompt})

        await _run_loop(current_context, new_messages, config, signal, stream, stream_fn)

    _spawn_runner(stream, runner())
    return stream


def agent_loop_continue(
    context: AgentContext | Mapping[str, Any],
    config: AgentLoopConfig | Mapping[str, Any],
    signal: Any | None = None,
    stream_fn: StreamFn | None = None,
) -> EventStream[AgentEvent, list[AgentMessage]]:
    messages = list(_get_context_messages(context))
    if not messages:
        raise ValueError("Cannot continue: no messages in context")

    if message_role(messages[-1]) == "assistant":
        raise ValueError("Cannot continue from message role: assistant")

    stream = _create_agent_stream()

    async def runner() -> None:
        new_messages: list[AgentMessage] = []
        current_context = _build_context(context)
        stream.push({"type": "agent_start"})
        stream.push({"type": "turn_start"})
        await _run_loop(current_context, new_messages, config, signal, stream, stream_fn)

    _spawn_runner(stream, runner())
    return stream


def _spawn_runner(stream: EventStream[AgentEvent, list[AgentMessage]], coro: asyncio.Future[Any] | Any) -> None:
    try:
        loop = asyncio.get_running_loop()
    except RuntimeError as exc:
        raise RuntimeError("agent_loop requires an active asyncio event loop") from exc

    task = loop.create_task(cast("asyncio.coroutines.Coroutine[Any, Any, None]", coro))

    def done_callback(done_task: asyncio.Task[Any]) -> None:
        if done_task.cancelled():
            stream.end([])
            return
        exc = done_task.exception()
        if exc is not None:
            stream.set_exception(exc)

    task.add_done_callback(done_callback)
    setattr(stream, "_task", task)


def _create_agent_stream() -> EventStream[AgentEvent, list[AgentMessage]]:
    return EventStream(
        lambda event: cast(str, event.get("type", "")) == "agent_end",
        lambda event: cast(list[AgentMessage], event.get("messages", [])),
    )


def _get_context_messages(context: AgentContext | Mapping[str, Any]) -> Sequence[AgentMessage]:
    messages = get_value(context, "messages", default=[])
    if isinstance(messages, Sequence):
        return cast(Sequence[AgentMessage], messages)
    return []


def _build_context(
    context: AgentContext | Mapping[str, Any],
    append: Sequence[AgentMessage] | None = None,
) -> AgentContext:
    messages = list(_get_context_messages(context))
    if append:
        messages.extend(append)
    system_prompt = get_value(context, "systemPrompt", "system_prompt", default="")
    tools = get_value(context, "tools", default=None)
    built: AgentContext = {"systemPrompt": cast(str, system_prompt), "messages": messages}
    if isinstance(tools, list):
        built["tools"] = cast(list[AgentTool], tools)
    return built


async def _run_loop(
    current_context: AgentContext,
    new_messages: list[AgentMessage],
    config: AgentLoopConfig | Mapping[str, Any],
    signal: Any | None,
    stream: EventStream[AgentEvent, list[AgentMessage]],
    stream_fn: StreamFn | None = None,
) -> None:
    first_turn = True
    pending_messages = await _get_queued_messages(config, "getSteeringMessages", "get_steering_messages")

    while True:
        has_more_tool_calls = True
        steering_after_tools: list[AgentMessage] | None = None

        while has_more_tool_calls or pending_messages:
            if first_turn:
                first_turn = False
            else:
                stream.push({"type": "turn_start"})

            if pending_messages:
                for message in pending_messages:
                    stream.push({"type": "message_start", "message": message})
                    stream.push({"type": "message_end", "message": message})
                    current_context["messages"].append(message)
                    new_messages.append(message)
                pending_messages = []

            message = await _stream_assistant_response(current_context, config, signal, stream, stream_fn)
            new_messages.append(message)

            reason = stop_reason(message)
            if reason in {"error", "aborted"}:
                stream.push({"type": "turn_end", "message": message, "toolResults": []})
                stream.push({"type": "agent_end", "messages": new_messages})
                stream.end(new_messages)
                return

            tool_calls = _extract_tool_calls(message)
            has_more_tool_calls = len(tool_calls) > 0

            tool_results: list[ToolResultMessage] = []
            if has_more_tool_calls:
                execution = await _execute_tool_calls(
                    cast(list[AgentTool] | None, current_context.get("tools")),
                    message,
                    signal,
                    stream,
                    cast(Any, get_value(config, "getSteeringMessages", "get_steering_messages", default=None)),
                )
                tool_results.extend(execution.tool_results)
                steering_after_tools = execution.steering_messages
                for result in tool_results:
                    current_context["messages"].append(result)
                    new_messages.append(result)

            stream.push({"type": "turn_end", "message": message, "toolResults": tool_results})

            if steering_after_tools:
                pending_messages = steering_after_tools
                steering_after_tools = None
            else:
                pending_messages = await _get_queued_messages(config, "getSteeringMessages", "get_steering_messages")

        follow_up_messages = await _get_queued_messages(config, "getFollowUpMessages", "get_follow_up_messages")
        if follow_up_messages:
            pending_messages = follow_up_messages
            continue

        break

    stream.push({"type": "agent_end", "messages": new_messages})
    stream.end(new_messages)


async def _stream_assistant_response(
    context: AgentContext,
    config: AgentLoopConfig | Mapping[str, Any],
    signal: Any | None,
    stream: EventStream[AgentEvent, list[AgentMessage]],
    stream_fn: StreamFn | None = None,
) -> AssistantMessage:
    messages: list[AgentMessage] = list(context["messages"])

    transform_context = get_value(config, "transformContext", "transform_context", default=None)
    if callable(transform_context):
        transformed = await maybe_await(transform_context(messages, signal))
        messages = list(transformed)

    convert_to_llm = get_value(config, "convertToLlm", "convert_to_llm", default=None)
    if not callable(convert_to_llm):
        raise ValueError("AgentLoopConfig.convertToLlm is required")

    llm_messages = await maybe_await(convert_to_llm(messages))
    llm_context: Context = {
        "systemPrompt": cast(str, context.get("systemPrompt", "")),
        "messages": cast(list[Message], list(llm_messages)),
    }
    if "tools" in context and context["tools"] is not None:
        llm_context["tools"] = cast(list[Any], context["tools"])

    stream_function = stream_fn or stream_simple
    model = cast(Model, get_value(config, "model", default=None))
    if not model:
        raise ValueError("AgentLoopConfig.model is required")

    get_api_key = get_value(config, "getApiKey", "get_api_key", default=None)
    resolved_api_key = None
    if callable(get_api_key):
        resolved_api_key = await maybe_await(get_api_key(model_provider(model)))
    if not resolved_api_key:
        resolved_api_key = get_value(config, "apiKey", "api_key", default=None)

    options = _config_to_options(config)
    options["apiKey"] = resolved_api_key
    options["signal"] = signal

    response = await maybe_await(stream_function(model, llm_context, options))

    partial_message: AssistantMessage | None = None
    added_partial = False

    async for event in response:
        event_type = get_value(event, "type", default="")
        if event_type == "start":
            partial_message = cast(AssistantMessage, get_value(event, "partial", default={}))
            context["messages"].append(partial_message)
            added_partial = True
            stream.push({"type": "message_start", "message": clone_message(partial_message)})
            continue

        if event_type in {
            "text_start",
            "text_delta",
            "text_end",
            "thinking_start",
            "thinking_delta",
            "thinking_end",
            "toolcall_start",
            "toolcall_delta",
            "toolcall_end",
        }:
            if partial_message is not None:
                partial_message = cast(AssistantMessage, get_value(event, "partial", default=partial_message))
                context["messages"][-1] = partial_message
                stream.push(
                    {
                        "type": "message_update",
                        "assistantMessageEvent": event,
                        "message": clone_message(partial_message),
                    }
                )
            continue

        if event_type in {"done", "error"}:
            final_message = cast(AssistantMessage, await response.result())
            if added_partial:
                context["messages"][-1] = final_message
            else:
                context["messages"].append(final_message)
                stream.push({"type": "message_start", "message": clone_message(final_message)})
            stream.push({"type": "message_end", "message": final_message})
            return final_message

    return cast(AssistantMessage, await response.result())


@dataclass(slots=True)
class _ToolExecutionResult:
    tool_results: list[ToolResultMessage]
    steering_messages: list[AgentMessage] | None = None


async def _execute_tool_calls(
    tools: list[AgentTool] | None,
    assistant_message: AssistantMessage,
    signal: Any | None,
    stream: EventStream[AgentEvent, list[AgentMessage]],
    get_steering_messages: Any = None,
) -> _ToolExecutionResult:
    tool_calls = _extract_tool_calls(assistant_message)
    results: list[ToolResultMessage] = []
    steering_messages: list[AgentMessage] | None = None

    for index, tool_call in enumerate(tool_calls):
        tool_name = cast(str, tool_call.get("name", ""))
        tool_id = cast(str, tool_call.get("id", ""))
        tool_args = tool_call.get("arguments", {})
        tool = _find_tool(tools, tool_name)

        stream.push(
            {
                "type": "tool_execution_start",
                "toolCallId": tool_id,
                "toolName": tool_name,
                "args": tool_args,
            }
        )

        result: AgentToolResult
        is_error = False
        try:
            if tool is None:
                raise ValueError(f"Tool {tool_name} not found")

            validated_args = validate_tool_arguments(tool, tool_call)
            result = await _execute_tool(tool, tool_id, validated_args, signal, stream, tool_name, tool_args)
        except Exception as exc:
            result = {
                "content": [{"type": "text", "text": str(exc)}],
                "details": {},
            }
            is_error = True

        stream.push(
            {
                "type": "tool_execution_end",
                "toolCallId": tool_id,
                "toolName": tool_name,
                "result": result,
                "isError": is_error,
            }
        )

        tool_result: ToolResultMessage = {
            "role": "toolResult",
            "toolCallId": tool_id,
            "toolName": tool_name,
            "content": cast(list[Any], result.get("content", [])),
            "details": result.get("details", {}),
            "isError": is_error,
            "timestamp": now_ms(),
        }
        results.append(tool_result)
        stream.push({"type": "message_start", "message": tool_result})
        stream.push({"type": "message_end", "message": tool_result})

        if callable(get_steering_messages):
            steering = await maybe_await(get_steering_messages())
            steering_list = list(steering or [])
            if steering_list:
                steering_messages = steering_list
                for skipped in tool_calls[index + 1 :]:
                    results.append(_skip_tool_call(skipped, stream))
                break

    return _ToolExecutionResult(tool_results=results, steering_messages=steering_messages)


async def _execute_tool(
    tool: AgentTool,
    tool_call_id: str,
    validated_args: Any,
    signal: Any | None,
    stream: EventStream[AgentEvent, list[AgentMessage]],
    tool_name: str,
    tool_args: Any,
) -> AgentToolResult:
    execute = get_value(tool, "execute", default=None)
    if not callable(execute):
        raise ValueError(f"Tool {tool_name} missing execute()")

    def on_update(partial_result: AgentToolResult) -> None:
        stream.push(
            {
                "type": "tool_execution_update",
                "toolCallId": tool_call_id,
                "toolName": tool_name,
                "args": tool_args,
                "partialResult": partial_result,
            }
        )

    result = await maybe_await(execute(tool_call_id, validated_args, signal, on_update))
    if not isinstance(result, Mapping):
        raise TypeError(f"Tool {tool_name} returned invalid result type: {type(result)!r}")
    normalized = dict(result)
    normalized.setdefault("content", [])
    normalized.setdefault("details", {})
    return cast(AgentToolResult, normalized)


def _skip_tool_call(
    tool_call: ToolCall,
    stream: EventStream[AgentEvent, list[AgentMessage]],
) -> ToolResultMessage:
    tool_name = cast(str, tool_call.get("name", ""))
    tool_id = cast(str, tool_call.get("id", ""))
    args = tool_call.get("arguments", {})
    result: AgentToolResult = {
        "content": [{"type": "text", "text": "Skipped due to queued user message."}],
        "details": {},
    }

    stream.push({"type": "tool_execution_start", "toolCallId": tool_id, "toolName": tool_name, "args": args})
    stream.push(
        {
            "type": "tool_execution_end",
            "toolCallId": tool_id,
            "toolName": tool_name,
            "result": result,
            "isError": True,
        }
    )

    tool_result: ToolResultMessage = {
        "role": "toolResult",
        "toolCallId": tool_id,
        "toolName": tool_name,
        "content": cast(list[Any], result["content"]),
        "details": {},
        "isError": True,
        "timestamp": now_ms(),
    }
    stream.push({"type": "message_start", "message": tool_result})
    stream.push({"type": "message_end", "message": tool_result})
    return tool_result


def _extract_tool_calls(message: AgentMessage) -> list[ToolCall]:
    content = message_content(message)
    if not isinstance(content, list):
        return []
    tool_calls: list[ToolCall] = []
    for block in content:
        if isinstance(block, Mapping) and block.get("type") == "toolCall":
            tool_calls.append(cast(ToolCall, block))
    return tool_calls


def _find_tool(tools: Sequence[AgentTool] | None, name: str) -> AgentTool | None:
    if tools is None:
        return None
    for tool in tools:
        if get_value(tool, "name", default="") == name:
            return tool
    return None


async def _get_queued_messages(
    config: AgentLoopConfig | Mapping[str, Any],
    *keys: str,
) -> list[AgentMessage]:
    getter = get_value(config, *keys, default=None)
    if not callable(getter):
        return []
    messages = await maybe_await(getter())
    return list(messages or [])


def _config_to_options(config: AgentLoopConfig | Mapping[str, Any]) -> dict[str, Any]:
    if isinstance(config, Mapping):
        return dict(config)

    result: dict[str, Any] = {}
    for key in dir(config):
        if key.startswith("_"):
            continue
        value = getattr(config, key, None)
        if callable(value):
            continue
        result[key] = value
    return result


class _ValidationError(ValueError):
    pass


def validate_tool_arguments(tool: AgentTool, tool_call: ToolCall) -> Any:
    schema = get_value(tool, "parameters", default={})
    arguments = copy.deepcopy(cast(dict[str, Any], tool_call.get("arguments", {})))
    if not isinstance(schema, Mapping):
        return arguments

    try:
        return _coerce_schema_value(schema, arguments, path="root")
    except _ValidationError as exc:
        tool_name = cast(str, tool_call.get("name", get_value(tool, "name", default="tool")))
        pretty_args = json.dumps(arguments, indent=2, sort_keys=True, default=str)
        raise ValueError(
            f'Validation failed for tool "{tool_name}":\n  - {exc}\n\nReceived arguments:\n{pretty_args}'
        ) from exc


def _coerce_schema_value(schema: Mapping[str, Any], value: Any, *, path: str) -> Any:
    if "const" in schema and value != schema["const"]:
        raise _ValidationError(f"{path}: expected constant {schema['const']!r}")

    if "enum" in schema:
        enum_values = schema["enum"]
        if isinstance(enum_values, list) and value not in enum_values:
            raise _ValidationError(f"{path}: expected one of {enum_values!r}")

    for composite_key in ("oneOf", "anyOf"):
        if composite_key in schema and isinstance(schema[composite_key], list):
            errors: list[str] = []
            for option in schema[composite_key]:
                if not isinstance(option, Mapping):
                    continue
                try:
                    return _coerce_schema_value(option, value, path=path)
                except _ValidationError as err:
                    errors.append(str(err))
            if errors:
                raise _ValidationError(f"{path}: no matching {composite_key} schema ({'; '.join(errors)})")

    schema_type = schema.get("type")
    if isinstance(schema_type, list):
        errors: list[str] = []
        for option_type in schema_type:
            option_schema = dict(schema)
            option_schema["type"] = option_type
            try:
                return _coerce_schema_value(option_schema, value, path=path)
            except _ValidationError as err:
                errors.append(str(err))
        raise _ValidationError(f"{path}: {', '.join(errors)}")

    if schema_type == "object" or ("properties" in schema and schema_type is None):
        if not isinstance(value, Mapping):
            raise _ValidationError(f"{path}: expected object")

        properties = schema.get("properties", {})
        if not isinstance(properties, Mapping):
            properties = {}
        required = schema.get("required", [])
        if isinstance(required, list):
            for field_name in required:
                if field_name not in value:
                    raise _ValidationError(f"{path}: missing required field {field_name!r}")

        additional_properties = schema.get("additionalProperties", True)
        result: dict[str, Any] = {}
        for key, raw_item in value.items():
            if key in properties and isinstance(properties[key], Mapping):
                result[key] = _coerce_schema_value(cast(Mapping[str, Any], properties[key]), raw_item, path=f"{path}.{key}")
                continue

            if additional_properties is False:
                raise _ValidationError(f"{path}: unexpected field {key!r}")
            if isinstance(additional_properties, Mapping):
                result[key] = _coerce_schema_value(additional_properties, raw_item, path=f"{path}.{key}")
            else:
                result[key] = raw_item

        return result

    if schema_type == "array":
        if not isinstance(value, list):
            raise _ValidationError(f"{path}: expected array")
        items_schema = schema.get("items")
        coerced_items = value
        if isinstance(items_schema, Mapping):
            coerced_items = [
                _coerce_schema_value(items_schema, item, path=f"{path}[{index}]") for index, item in enumerate(value)
            ]

        min_items = schema.get("minItems")
        max_items = schema.get("maxItems")
        if isinstance(min_items, int) and len(coerced_items) < min_items:
            raise _ValidationError(f"{path}: expected at least {min_items} items")
        if isinstance(max_items, int) and len(coerced_items) > max_items:
            raise _ValidationError(f"{path}: expected at most {max_items} items")
        return coerced_items

    if schema_type == "string":
        if value is None:
            raise _ValidationError(f"{path}: expected string")
        coerced = value if isinstance(value, str) else str(value)
        min_length = schema.get("minLength")
        max_length = schema.get("maxLength")
        if isinstance(min_length, int) and len(coerced) < min_length:
            raise _ValidationError(f"{path}: expected minimum length {min_length}")
        if isinstance(max_length, int) and len(coerced) > max_length:
            raise _ValidationError(f"{path}: expected maximum length {max_length}")
        return coerced

    if schema_type == "integer":
        coerced = _coerce_number(value, integer=True, path=path)
        _validate_numeric_bounds(schema, coerced, path=path)
        return coerced

    if schema_type == "number":
        coerced = _coerce_number(value, integer=False, path=path)
        _validate_numeric_bounds(schema, coerced, path=path)
        return coerced

    if schema_type == "boolean":
        if isinstance(value, bool):
            return value
        if isinstance(value, int) and value in {0, 1}:
            return bool(value)
        if isinstance(value, str):
            lowered = value.strip().lower()
            if lowered in {"true", "1"}:
                return True
            if lowered in {"false", "0"}:
                return False
        raise _ValidationError(f"{path}: expected boolean")

    if schema_type == "null":
        if value is None:
            return None
        raise _ValidationError(f"{path}: expected null")

    return value


def _coerce_number(value: Any, *, integer: bool, path: str) -> int | float:
    if isinstance(value, bool):
        raise _ValidationError(f"{path}: expected {'integer' if integer else 'number'}")

    if integer:
        if isinstance(value, int):
            return value
        if isinstance(value, float) and value.is_integer():
            return int(value)
        if isinstance(value, str):
            stripped = value.strip()
            try:
                parsed = float(stripped)
            except ValueError as exc:
                raise _ValidationError(f"{path}: expected integer") from exc
            if not parsed.is_integer():
                raise _ValidationError(f"{path}: expected integer")
            return int(parsed)
        raise _ValidationError(f"{path}: expected integer")

    if isinstance(value, (int, float)) and not isinstance(value, bool):
        return float(value)
    if isinstance(value, str):
        stripped = value.strip()
        try:
            return float(stripped)
        except ValueError as exc:
            raise _ValidationError(f"{path}: expected number") from exc
    raise _ValidationError(f"{path}: expected number")


def _validate_numeric_bounds(schema: Mapping[str, Any], value: int | float, *, path: str) -> None:
    minimum = schema.get("minimum")
    maximum = schema.get("maximum")
    exclusive_minimum = schema.get("exclusiveMinimum")
    exclusive_maximum = schema.get("exclusiveMaximum")
    if isinstance(minimum, (int, float)) and value < minimum:
        raise _ValidationError(f"{path}: expected >= {minimum}")
    if isinstance(maximum, (int, float)) and value > maximum:
        raise _ValidationError(f"{path}: expected <= {maximum}")
    if isinstance(exclusive_minimum, (int, float)) and value <= exclusive_minimum:
        raise _ValidationError(f"{path}: expected > {exclusive_minimum}")
    if isinstance(exclusive_maximum, (int, float)) and value >= exclusive_maximum:
        raise _ValidationError(f"{path}: expected < {exclusive_maximum}")


def agentLoop(
    prompts: Sequence[AgentMessage],
    context: AgentContext | Mapping[str, Any],
    config: AgentLoopConfig | Mapping[str, Any],
    signal: Any | None = None,
    stream_fn: StreamFn | None = None,
) -> EventStream[AgentEvent, list[AgentMessage]]:
    return agent_loop(prompts, context, config, signal=signal, stream_fn=stream_fn)


def agentLoopContinue(
    context: AgentContext | Mapping[str, Any],
    config: AgentLoopConfig | Mapping[str, Any],
    signal: Any | None = None,
    stream_fn: StreamFn | None = None,
) -> EventStream[AgentEvent, list[AgentMessage]]:
    return agent_loop_continue(context, config, signal=signal, stream_fn=stream_fn)
