from __future__ import annotations

import time
import uuid
from typing import Iterator, AsyncIterator

from turbo_agent_core.schema.enums import JSON
from turbo_agent_core.schema.events import (
    BaseEvent,
    RunLifecycleCreatedEvent,
    RunLifecycleCreatedPayload,
    RunLifecycleCompletedEvent,
    RunLifecycleCompletedPayload,
    RunLifecycleFailedEvent,
    RunLifecycleFailedPayload,
    ContentAgentResultStartEvent,
    ContentAgentResultStartPayload,
    ContentAgentResultDeltaEvent,
    ContentAgentResultDeltaPayload,
    ContentAgentResultEndEvent,
    ContentAgentResultEndPayload,
    ControlRollbackEvent,
    ControlRollbackPayload,
    ExecutorMetadata,
    UserInfo,
)
from turbo_agent_core.schema.agents import LLMTool
from turbo_agent_core.utils.json_stream import JsonStreamParser
from turbo_agent_runtime.utils.file_materializer import materialize_knowledge_resources
from .base import construct_json_correction_conversation, construct_tool_messages, parse_json_from_llm_output
from turbo_agent_runtime.utils.executor_identity import build_executor_id


class LLMToolRuntime(LLMTool):
    """运行 LLMTool 版本：绑定模型并调用；可流式。"""

    def _resolve_user_metadata(self, **kwargs):
        raw = kwargs.get("user_metadata") or kwargs.get("user_info")
        if isinstance(raw, UserInfo):
            return raw
        if isinstance(raw, dict):
            try:
                return UserInfo.model_validate(raw)
            except Exception:
                pass
        user_id = kwargs.get("user_id")
        username = kwargs.get("username")
        return UserInfo(id=str(user_id or "local"), username=str(username or "local"))

    def run(self, input: JSON, **kwargs) -> JSON:
        data = input if isinstance(input, dict) else {"text": str(input)}
        try:
            valid_input = self.validate_input(data)
            schema = getattr(self, "input_schema", {}) or {}
            # if isinstance(schema, dict) and schema:
            #     valid_input = materialize_knowledge_resources(schema, valid_input)
        except ValueError as e:
            return {"error": str(e), "tool_id": self.id}

        conversation, leaf_id = construct_tool_messages(
            tool_id=self.id,
            tool_name=getattr(self, "name", self.id),
            input_params=getattr(self, "input", []),
            output_params=getattr(self, "output", []),
            setting=getattr(self, "setting", None),
            input_data=valid_input
        )

        response_message = self.model.run(conversation, leaf_message_id=leaf_id)
        response_text = str(response_message.content)

        try:
            final_result = parse_json_from_llm_output(response_text)
            shaped = self.validate_output(final_result)
            return {"tool_id": self.id, "version_id": getattr(self, "id", None), "output": shaped, "raw": final_result}
        except Exception as e:
            try:
                output_schema = getattr(self, "output_schema", {})
                conversation, leaf_id = construct_json_correction_conversation(self.id, output_schema, response_text, str(e))
                response_message = self.model.run(conversation, leaf_id)
                corrected_text = str(response_message.content)
                final_result = parse_json_from_llm_output(corrected_text)
                shaped = self.validate_output(final_result)
                return {"tool_id": self.id, "version_id": getattr(self, "id", None), "output": shaped, "raw": final_result}
            except Exception as e2:
                return {"error": f"Execution failed after retry: {str(e2)}", "tool_id": self.id}

    def stream(self, input: JSON, **kwargs) -> Iterator[BaseEvent]:
        start_time = time.time()
        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]

        data = input if isinstance(input, dict) else {"text": str(input)}

        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=self._resolve_user_metadata(**kwargs),
            payload=RunLifecycleCreatedPayload(input_data=data),
        )
        try:
            valid_input = self.validate_input(data)
            schema = getattr(self, "input_schema", {}) or {}
            # if isinstance(schema, dict) and schema:
            #     valid_input = materialize_knowledge_resources(schema, valid_input)
        except ValueError as e:
            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": "ValidationError: Input parse illegal", "message": str(e)}),
            )
            return

        conversation, leaf_id = construct_tool_messages(
            tool_id=self.id,
            tool_name=getattr(self, "name", self.id),
            input_params=getattr(self, "input", []),
            output_params=getattr(self, "output", []),
            setting=getattr(self, "setting", None),
            input_data=valid_input
        )

        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="json")
        )

        parser = JsonStreamParser()
        full_text = ""
        start_timestamp = time.time() * 1000

        collected_usage = {
            "promptToken": 0,
            "completionToken": 0,
            "reasoningToken": 0,
            "tokenCost": 0,
            "moneyCost": 0.0
        }
        reasoning_full = ""

        try:
            for event in self.model.stream(conversation, leaf_message_id=leaf_id):
                if event.type == "run.lifecycle.completed" and event.payload.usage:
                    u = event.payload.usage
                    collected_usage["promptToken"] += u.get("promptToken", 0)
                    collected_usage["completionToken"] += u.get("completionToken", 0)
                    collected_usage["tokenCost"] += u.get("tokenCost", 0)
                    collected_usage["reasoningToken"] += u.get("reasoningToken", 0)
                    collected_usage["moneyCost"] += u.get("moneyCost", 0.0)
                    
                if event.type.startswith("content.reasoning"):
                    if event.type == "content.reasoning.start":
                        yield ContentAgentResultDeltaEvent(
                            trace_id=trace_id, run_id=run_id,
                            trace_path=[],
                            executor_type=executor_type,
                            executor_id=executor_id, executor_path=executor_path,
                            payload=ContentAgentResultDeltaPayload(delta="", key_path=path, part="reasoning")
                        )
                    elif event.type == "content.reasoning.delta":
                        yield ContentAgentResultDeltaEvent(
                            trace_id=trace_id, run_id=run_id,
                            trace_path=[],
                            executor_type=executor_type,
                            executor_id=executor_id, executor_path=executor_path,
                            payload=ContentAgentResultDeltaPayload(delta=event.payload.delta, key_path=path, part="reasoning")
                        )
                    elif event.type == "content.reasoning.end":
                        reasoning_full = event.payload.full_text

                if event.type != "content.text.delta":
                    continue
                chunk = event.payload.delta
                if not chunk:
                    continue

                full_text += chunk
                deltas = parser.feed(chunk)
                for path, value in deltas:
                    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=value, key_path=path, part="output")
                    )

            try:
                final_result = parse_json_from_llm_output(full_text)
                shaped = self.validate_output(final_result)
            except Exception as e:
                yield ControlRollbackEvent(
                    trace_id=trace_id, run_id=run_id,
                    executor_type=executor_type,
                    executor_id=executor_id, executor_path=executor_path,
                    payload=ControlRollbackPayload(target_timestamp=start_timestamp, reason="ValidationError: Output format illegal,retrying JSON correction")
                )
                output_schema = getattr(self, "output_schema", {})
                conversation, leaf_id = construct_json_correction_conversation(self.id, output_schema, full_text, str(e))
                parser = JsonStreamParser()
                corrected_json_str = ""
                for event in self.model.stream(conversation, leaf_id):
                    if event.type == "run.lifecycle.completed" and event.payload.usage:
                        u = event.payload.usage
                        collected_usage["promptToken"] += u.get("promptToken", 0)
                        collected_usage["completionToken"] += u.get("completionToken", 0)
                        collected_usage["tokenCost"] += u.get("tokenCost", 0)
                        collected_usage["moneyCost"] += u.get("moneyCost", 0.0)
                        collected_usage["reasoningToken"] += u.get("reasoningToken", 0)

                    if event.type == "content.text.delta":
                        chunk = event.payload.delta
                        corrected_json_str += chunk
                        deltas = parser.feed(chunk)
                        for path, value in deltas:
                            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=value, key_path=path, part="output")
                            )
                full_text = corrected_json_str
                final_result = parse_json_from_llm_output(full_text)
                shaped = self.validate_output(final_result)

            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=shaped,full_reasoning=reasoning_full)
            )
            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=None, usage={
                    **collected_usage,
                    "timeCost": time.time() - start_time
                }),
            )
        except Exception as e:
            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="error",full_result=str(e), full_reasoning=reasoning_full)
            )
            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)}),
            )

    async def a_run(self, input: JSON, **kwargs) -> JSON:
        data = input if isinstance(input, dict) else {"text": str(input)}
        try:
            valid_input = self.validate_input(data)
            schema = getattr(self, "input_schema", {}) or {}
            # if isinstance(schema, dict) and schema:
            #     valid_input = materialize_knowledge_resources(schema, valid_input)
        except ValueError as e:
            return {"error": str(e), "tool_id": self.id}

        conversation, leaf_id = construct_tool_messages(
            tool_id=self.id,
            tool_name=getattr(self, "name", self.id),
            input_params=getattr(self, "input", []),
            output_params=getattr(self, "output", []),
            setting=getattr(self, "setting", None),
            input_data=valid_input
        )

        response_message = await self.model.a_run(conversation, leaf_id)
        response_text = str(response_message.content)

        try:
            final_result = parse_json_from_llm_output(response_text)
            shaped = self.validate_output(final_result)
            return {"tool_id": self.id, "version_id": getattr(self, "id", None), "output": shaped, "raw": final_result}
        except Exception as e:
            return {"error": f"Execution failed: {str(e)}", "tool_id": self.id}

    async def a_stream(self, input: JSON, **kwargs) -> AsyncIterator[BaseEvent]:
        start_time = time.time()
        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]

        data = input if isinstance(input, dict) else {"text": str(input)}
        

        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=self._resolve_user_metadata(**kwargs),
            payload=RunLifecycleCreatedPayload(input_data=data),
        )
        try:
            valid_input = self.validate_input(data)
            schema = getattr(self, "input_schema", {}) or {}
            # if isinstance(schema, dict) and schema:
            #     valid_input = materialize_knowledge_resources(schema, valid_input)
        except ValueError as e:
            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": "ValidationError: Input format illegal", "message": str(e)}),
            )
            return

        conversation, leaf_id = construct_tool_messages(
            tool_id=self.id,
            tool_name=getattr(self, "name", self.id),
            input_params=getattr(self, "input", []),
            output_params=getattr(self, "output", []),
            setting=getattr(self, "setting", None),
            input_data=valid_input
        )

        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="json")
        )

        parser = JsonStreamParser()
        full_text = ""
        reasoning_full = ""
        start_timestamp = time.time() * 1000

        collected_usage = {
            "promptToken": 0,
            "reasoningToken": 0,
            "completionToken": 0,
            "tokenCost": 0,
            "moneyCost": 0.0
        }

        try:
            async for event in self.model.a_stream(conversation, leaf_message_id=leaf_id):
                if event.type == "run.lifecycle.completed" and event.payload.usage:
                    u = event.payload.usage
                    collected_usage["promptToken"] += u.get("promptToken", 0)
                    collected_usage["completionToken"] += u.get("completionToken", 0)
                    collected_usage["reasoningToken"] += u.get("reasoningToken", 0)
                    collected_usage["tokenCost"] += u.get("tokenCost", 0)
                    collected_usage["moneyCost"] += u.get("moneyCost", 0.0)
                
                if event.type.startswith("content.reasoning"):
                    if event.type == "content.reasoning.start":
                        yield ContentAgentResultDeltaEvent(
                            trace_id=trace_id, run_id=run_id,
                            trace_path=[],
                            executor_type=executor_type,
                            executor_id=executor_id, executor_path=executor_path,
                            payload=ContentAgentResultDeltaPayload(delta="", key_path=None, part="reasoning")
                        )
                    elif event.type == "content.reasoning.delta":
                        yield ContentAgentResultDeltaEvent(
                            trace_id=trace_id, run_id=run_id,
                            trace_path=[],
                            executor_type=executor_type,
                            executor_id=executor_id, executor_path=executor_path,
                            payload=ContentAgentResultDeltaPayload(delta=event.payload.delta, key_path=None, part="reasoning")
                        )
                    elif event.type == "content.reasoning.end":
                        reasoning_full = event.payload.full_text

                if event.type != "content.text.delta":
                    continue
                chunk = event.payload.delta
                if not chunk:
                    continue
                full_text += chunk
                deltas = parser.feed(chunk)
                for path, value in deltas:
                    yield ContentAgentResultDeltaEvent(
                        trace_id=trace_id, run_id=run_id,
                        trace_path=[],
                        executor_id=executor_id, executor_path=executor_path,
                        executor_type=executor_type,
                        payload=ContentAgentResultDeltaPayload(delta=value, key_path=path, part="output")
                    )

            try:
                final_result = parse_json_from_llm_output(full_text)
                shaped = self.validate_output(final_result)
            except Exception as e:
                yield ControlRollbackEvent(
                    trace_id=trace_id, run_id=run_id,
                    trace_path=[],
                    executor_id=executor_id, executor_path=executor_path,
                    executor_type=executor_type,
                    payload=ControlRollbackPayload(target_timestamp=start_timestamp, reason="ValidationError: Output format illegal,retrying JSON correction")
                )
                output_schema = getattr(self, "output_schema", {})
                conversation, leaf_id = construct_json_correction_conversation(self.id, output_schema, full_text, str(e))
                parser = JsonStreamParser()
                corrected_json_str = ""
                async for event in self.model.a_stream(conversation, leaf_message_id=leaf_id):
                    if event.type == "run.lifecycle.completed" and event.payload.usage:
                        u = event.payload.usage
                        collected_usage["promptToken"] += u.get("promptToken", 0)
                        collected_usage["completionToken"] += u.get("completionToken", 0)
                        collected_usage["reasoningToken"] += u.get("reasoningToken", 0)
                        collected_usage["tokenCost"] += u.get("tokenCost", 0)
                        collected_usage["moneyCost"] += u.get("moneyCost", 0.0)

                    if event.type == "content.text.delta":
                        chunk = event.payload.delta
                        corrected_json_str += chunk
                        deltas = parser.feed(chunk)
                        for path, value in deltas:
                            yield ContentAgentResultDeltaEvent(
                                trace_id=trace_id, run_id=run_id,
                                trace_path=[],
                                executor_id=executor_id, executor_path=executor_path,
                                executor_type=executor_type,
                                payload=ContentAgentResultDeltaPayload(delta=value, key_path=path, part="output")
                            )
                full_text = corrected_json_str
                final_result = parse_json_from_llm_output(full_text)
                shaped = self.validate_output(final_result)

            yield ContentAgentResultEndEvent(
                trace_id=trace_id, run_id=run_id,
                trace_path=[],
                executor_id=executor_id, executor_path=executor_path,
                executor_type=executor_type,
                payload=ContentAgentResultEndPayload(status="success", full_result=shaped,full_reasoning=reasoning_full)
            )
            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=None, usage={
                    **collected_usage,
                    "timeCost": time.time() - start_time
                }),
            )
        except Exception as e:
            yield ContentAgentResultEndEvent(
                trace_id=trace_id, run_id=run_id,
                trace_path=[],
                executor_id=executor_id, executor_path=executor_path,
                executor_type=executor_type,
                payload=ContentAgentResultEndPayload(status="error", full_result=str(e), full_reasoning=reasoning_full)
            )
            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)}),
            )





