from __future__ import annotations

import time
import uuid
from typing import AsyncIterator
from loguru import logger

from turbo_agent_core.schema.enums import JSON
from turbo_agent_core.schema.events import (
    BaseEvent,
    RunLifecycleCreatedEvent,
    RunLifecycleCreatedPayload,
    RunLifecycleCompletedEvent,
    RunLifecycleCompletedPayload,
    RunLifecycleFailedEvent,
    RunLifecycleFailedPayload,
    ContentTextDeltaEvent,
    ContentTextDeltaPayload,
    ContentAgentResultStartEvent,
    ContentAgentResultStartPayload,
    ContentAgentResultDeltaEvent,
    ContentAgentResultDeltaPayload,
    ContentAgentResultEndEvent,
    ContentAgentResultEndPayload,
    ExecutorMetadata,
    UserInfo,
)
from turbo_agent_core.schema.agents import AgentTool
from turbo_agent_core.utils.event_integrator import ChildEventStreamIntegrator
from turbo_agent_runtime.utils.file_materializer import materialize_knowledge_resources
from turbo_agent_runtime.utils.executor_identity import build_executor_id

class AgentToolRuntime(AgentTool):
    """运行 AgentTool 版本：委托给 backendAgent 的 runtime。"""

    async def run(self, input: JSON, **kwargs) -> JSON:
        ctx = kwargs.get("ctx")
        if not self.backendAgent:
            return {"error": "Missing backendAgent", "tool_id": self.id}
            
        data = input if isinstance(input, dict) else {"value": input}
        try:
            valid = self.validate_input(data)
            schema = getattr(self, "input_schema", {}) or {}
            if isinstance(schema, dict) and schema:
                valid = materialize_knowledge_resources(schema, valid)
        except ValueError as e:
            return {"error": str(e), "tool_id": self.id}
            
        # backendAgent.run returns Message
        raw_message = await self.backendAgent.run(valid)
        
        # Extract content and validate
        content = raw_message.content
        # Agent output is usually string. validate_output handles raw input.
        shaped = self.validate_output(content)
        
        return {"tool_id": self.id, "version_id": getattr(self, "id", None), "output": shaped, "raw": content}

    async def stream(self, input: JSON, **kwargs) -> AsyncIterator[BaseEvent]:
        trace_id = str(uuid.uuid4())
        run_id = f"run_{self.id}_{int(time.time()*1000)}"
        executor_type = self.run_type
        executor_id = build_executor_id(self, executor_type)
        executor_path = [executor_id]
        user_id = kwargs.get("user_id") or kwargs.get("userId") or "local"
        username = kwargs.get("username") or "local"

        yield RunLifecycleCreatedEvent(
            trace_id=trace_id,
            trace_path=[],
            run_id=run_id,
            run_path=[run_id],
            executor_type=executor_type,
            executor_id=executor_id,
            executor_path=executor_path,
            executor_metadata=ExecutorMetadata(
                id=self.id,
                name=self.name,
                belong_to_project_path=self.belong_to_project_path,
                name_id=self.name_id,
                description=self.description,
                avatar_uri=self.avatar_uri,
                run_type=self.run_type,
                version_id=self.version_id,
                version=self.version_tag,
            ),
            user_metadata=UserInfo(id=str(user_id), username=str(username)),
            payload=RunLifecycleCreatedPayload(input_data=input),
        )
        try:
            if not self.backendAgent:
                yield ContentTextDeltaEvent(
                    trace_id=trace_id, run_id=run_id,
                    executor_type=executor_type,
                    executor_id=executor_id, executor_path=executor_path,
                    payload=ContentTextDeltaPayload(delta="(无 backendAgent，无法执行)")
                )
                yield RunLifecycleFailedEvent(
                    trace_id=trace_id,
                    trace_path=[],
                    run_id=run_id,
                    run_path=[run_id],
                    executor_type=executor_type,
                    executor_id=executor_id,
                    executor_path=executor_path,
                    payload=RunLifecycleFailedPayload(error={"code": "ConfigError", "message": "Missing backendAgent"}),
                )
                return
            # 仅在流式前做一次资源实体化（假设输入静态）
            materialized_input = input
            schema = getattr(self, "input_schema", {}) or {}
            if isinstance(schema, dict) and schema:
                try:
                    materialized_input = materialize_knowledge_resources(schema, materialized_input)
                except Exception:
                    pass

            # Emit Result Start
            yield ContentAgentResultStartEvent(
                trace_id=trace_id, run_id=run_id,
                executor_type=executor_type,
                executor_id=executor_id, executor_path=executor_path,
                payload=ContentAgentResultStartPayload(mode="text")
            )

            # backendAgent 独立生成自己的 trace_id；这里按协议对其事件流做转发修补。
            integrator = ChildEventStreamIntegrator(
                root_trace_id=trace_id,
                parent_trace_path=[],
                parent_executor_path=executor_path,
                parent_run_id=run_id,
            )

            full_content = ""
            async for evt in integrator.integrate_async(self.backendAgent.a_stream(materialized_input)):
                if evt.type == "content.text.delta":
                    delta = evt.payload.delta
                    if delta:
                        full_content += delta
                        yield ContentAgentResultDeltaEvent(
                            trace_id=trace_id, run_id=run_id,
                            executor_type=executor_type, executor_id=executor_id, executor_path=executor_path,
                            payload=ContentAgentResultDeltaPayload(delta=delta, part="output")
                        )
                yield evt

            final_result_struct = {
               "tool_id": self.id, 
               "version_id": getattr(self, "id", None), 
               "output": self.validate_output(full_content), 
               "raw": full_content
            }

            yield ContentAgentResultEndEvent(
                trace_id=trace_id, run_id=run_id,
                executor_type=executor_type,
                executor_id=executor_id, executor_path=executor_path,
                payload=ContentAgentResultEndPayload(status="success", full_result=final_result_struct)
            )

            yield RunLifecycleCompletedEvent(
                trace_id=trace_id,
                trace_path=[],
                run_id=run_id,
                run_path=[run_id],
                executor_type=executor_type,
                executor_id=executor_id,
                executor_path=executor_path,
                payload=RunLifecycleCompletedPayload(output=final_result_struct, usage={"total_tokens": 0}),
            )
        except Exception as e:  # pragma: no cover
            yield RunLifecycleFailedEvent(
                trace_id=trace_id,
                trace_path=[],
                run_id=run_id,
                run_path=[run_id],
                executor_type=executor_type,
                executor_id=executor_id,
                executor_path=executor_path,
                payload=RunLifecycleFailedPayload(error={"code": "RuntimeError", "message": str(e)}),
            )
