# -*- coding: utf-8 -*-
import functools
import os
import threading

from typing import List, Dict, AsyncGenerator

from openai import OpenAI

from agentscope.agent import ReActAgent
from agentscope.formatter import DashScopeChatFormatter
from agentscope.message import TextBlock
from agentscope.model import DashScopeChatModel
from agentscope.pipeline import stream_printing_messages
from agentscope.tool import Toolkit, ToolResponse

from prompts import SYSTEM_PROMPT

from agentscope_runtime.adapters.agentscope.memory import (
    AgentScopeSessionHistoryMemory,
)
from agentscope_runtime.engine import AgentApp
from agentscope_runtime.engine.schemas.agent_schemas import (
    AgentRequest,
)
from agentscope_runtime.engine.services.agent_state.state_service import (
    InMemoryStateService,
)
from agentscope_runtime.engine.services.sandbox.sandbox_service import (
    SandboxService,
)

# flake8: noqa: E501
from agentscope_runtime.engine.services.session_history.session_history_service import (  # pylint: disable=line-too-long
    InMemorySessionHistoryService,  # pylint: disable=line-too-long
)


if os.path.exists(".env"):
    from dotenv import load_dotenv

    load_dotenv(".env")

USER_ID = "user_1"
SESSION_ID = "session_001"  # Using a fixed ID for simplicity

PORT = 8090


def sandbox_tool_adapter(func):
    """
    Sandbox Tool Adapter.
    Wraps a sandbox tool function so that its output is always converted
    into a ToolResponse object, which is required by the Toolkit.
    This adapter preserves the original function signature and docstring
    so that JSON schemas can be correctly generated by the Toolkit.
    Args:
        func: Original sandbox tool function (may return dict, string, etc.).
    Returns:
        A callable that produces ToolResponse instead of raw data.
    """

    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        res = func(*args, **kwargs)
        if isinstance(res, ToolResponse):
            return res
        # TODO: fix this
        return ToolResponse(content=[TextBlock(type="text", text=str(res))])

    return wrapper


desktop_url = ""


def init():
    agent_app = AgentApp(
        app_name="Friday",
        app_description="A helpful assistant",
    )

    @agent_app.init
    async def init_func(self):
        self.state_service = InMemoryStateService()
        self.session_service = InMemorySessionHistoryService()
        self.sandbox_service = SandboxService()

        await self.state_service.start()
        await self.session_service.start()
        await self.sandbox_service.start()

    @agent_app.shutdown
    async def shutdown_func(self):
        await self.state_service.stop()
        await self.session_service.stop()
        await self.sandbox_service.stop()

    @agent_app.query(framework="agentscope")
    async def query_func(
        self,
        msgs,
        request: AgentRequest = None,
    ):
        session_id = request.session_id
        user_id = request.user_id

        state = await self.state_service.export_state(
            session_id=session_id,
            user_id=user_id,
        )
        # Get sandbox
        sandboxes = self.sandbox_service.connect(
            session_id=session_id,
            user_id=user_id,
            sandbox_types=["browser"],
        )

        sandbox = sandboxes[0]
        global desktop_url
        desktop_url = sandbox.desktop_url

        browser_tools = [
            sandbox.browser_navigate,
            sandbox.browser_take_screenshot,
            sandbox.browser_snapshot,
            sandbox.browser_click,
            sandbox.browser_type,
        ]

        toolkit = Toolkit()
        for tool in browser_tools:
            toolkit.register_tool_function(sandbox_tool_adapter(tool))

        agent = ReActAgent(
            name="Friday",
            model=DashScopeChatModel(
                "qwen-max",
                api_key=os.getenv("DASHSCOPE_API_KEY"),
                enable_thinking=True,
                stream=True,
            ),
            sys_prompt=SYSTEM_PROMPT,
            toolkit=toolkit,
            memory=AgentScopeSessionHistoryMemory(
                service=self.session_service,
                session_id=session_id,
                user_id=user_id,
            ),
            formatter=DashScopeChatFormatter(),
        )

        if state:
            agent.load_state_dict(state)

        async for msg, last in stream_printing_messages(
            agents=[agent],
            coroutine_task=agent(msgs),
        ):
            yield msg, last

        state = agent.state_dict()

        await self.state_service.save_state(
            user_id=user_id,
            session_id=session_id,
            state=state,
        )

    def run_agent_app():
        agent_app.run(host="127.0.0.1", port=PORT)

    threading.Thread(target=run_agent_app, daemon=True).start()
    return agent_app


class AgentscopeBrowseruseAgent:
    def __init__(self) -> None:
        self.agent = init()
        self.desktop_url = ""

    async def chat(
        self,
        chat_messages: List[Dict],
    ) -> AsyncGenerator[Dict, None]:
        client = OpenAI(base_url=f"http://127.0.0.1:{PORT}/compatible-mode/v1")

        stream = client.responses.create(
            model="any_name",
            input=chat_messages[-1]["content"],
            stream=True,
        )
        global desktop_url

        self.desktop_url = desktop_url
        for chunk in stream:
            if hasattr(chunk, "delta"):
                yield chunk.delta
            else:
                yield {}
            # if chunk.choices[0].delta.content is not None:
            #    yield  chunk.choices[0].delta.content

    async def close(self) -> None:
        await self.sandbox_service.stop()
        await self.mem_service.stop()
