import re
import inspect
import time
import asyncio
import traceback

from temporary import generate_tmp_id
from stream import Stream

# Tool related classes must be moved to a different module

class AgentField:
    def __init__(self, name, tp, default=None):
        self.name=name
        self.type=tp
        self.default=default

# Instances are created and destroyed constantly, agents must not depend solely on in-memory data
# to be useful and consistent
class Agent:
    # Class variables, must be defined by subclasses of Agent:
    # - identifier (for kind in creation)
    # - schema: list[AgentField{name, type: str}] (for configuration)

    def __init__(self, app, data, resources=[]):
        self.app = app

        # These variables are set by '_set_basic_info'
        self.id = None
        self.name = None

        # Auto populated based on the schema fields
        agent_class = type(self)
        self.config = {field.name: data.get(field.name, field.default) for field in getattr(agent_class, 'schema_fields', [])}
        self.resources = resources

        # should be in expert init?
        self.tools = []
        self.remotes = []
        for resource in self.resources:
            if resource.kind == 'tool':
                if resource.value.get('type') == 'native' and 'name' in resource.value:
                    tool_cls = next((t for t in self.app.tool_classes if t.tool_name() == resource.value['name']), None)
                    if tool_cls:
                        tool = tool_cls()
                        spec = tool.get_spec().clone()
                        spec.name = resource.key
                        self.tools.append(spec)
            elif resource.kind == 'remote':
                value = resource.value or {}
                if isinstance(value, str):
                    value = {'prefix': value}
                self.remotes.append({
                    'name': resource.key,
                    'prefix': value.get('prefix'),
                })

    
    def _set_basic_info(self, id, name, user_id):
        self.id = id
        self.name = name
        self.user_id = user_id

    async def run(self, agent_run):
        """Core method. Subclasses override this to handle the AgentRun."""
        agent_run.close()
        raise NotImplementedError("Subclasses must implement run()")

    async def prompt(self, instruction, input=None, **kwargs):
        run = AgentRun.create(
            self.app,
            agent_id=self.id,
            user_id=self.user_id,
            input={
                "instruction": instruction,
                "input": input,
                **kwargs,
            },
        )
        out_stream = Stream(self.app)

        def on_output_event(stream):
            def on_chunk(chunk):
                out_stream.feed(chunk)
            stream.listen(on_chunk, lambda: None)
        def on_done(data):
            out_stream.close()
        run.on('output', on_output_event)
        run.on('done', on_done)

        # Change the output shape based on the presence of tools
        if 'tools' in kwargs:
            tool_calls = []
            tool_names = [t['function']['name'] for t in kwargs['tools'] if isinstance(t, dict)]
            def on_tool_call(call):
                print('prompt call', call)
                if call.tool_name in tool_names:
                    tool_calls.append(call)
            run.on('tool-call', on_tool_call)
            await self.run(run)

            # Cannot stream tool calls out, wait for them
            await out_stream.content_wait()
            print('tool names' + repr(tool_names))
            print('tool calls' + repr(tool_calls))
            return {
                'output': out_stream,
                'tool_calls': tool_calls,
            }
        else:
            await self.run(run)
            return out_stream


    async def received_message(self, message):
        chat = await self.app.get_chat(message.chat_id)
        thread = chat.get_thread(message.thread_id)
        thread.skip_list += ['log', 'resource']
        await thread.load_page()

        messages = []
        for msg in thread.messages:
            role = self._role_for_message(msg)
            messages.append({"role": role, "content": msg.content})

        remotes = []
        def add_remote(name):
            if name not in remotes and (self.user_id, name) in self.app.remotes:
                remotes.append(name)

        remote_messages = await thread.fetch_kind_page('resource', 'remote')
        for msg in remote_messages:
            add_remote(msg.content)

        run = AgentRun.create(
            self.app,
            agent_id=self.id,
            user_id=self.user_id,
            input={
                "messages": messages,
                "chat_id": message.chat_id,
                "thread_id": message.thread_id,
                "remotes": remotes,
            },
        )

        async def on_inference(msg):
            kwargs = {'content': msg.content}

            if msg.role == 'assistant':
                # kind and subkind are default
                pass
            elif msg.role == 'reasoning':
                kwargs['kind'] = "log"
                kwargs['subkind'] = "reasoning"
            elif msg.role == 'tool':
                kwargs['kind'] = "log"
                kwargs['subkind'] = "tool-call"
                kwargs['content'] = repr(msg.content)
            else:
                # Unknown message
                kwargs['kind'] = "log"

            new_message = await self.app.send_chat_message(
                sender=self.id,
                chat_id=message.chat_id,
                thread_id=message.thread_id,
                **kwargs
            )
            run.emit_event('chat-message', {
                'messageId': new_message['id'],
                'chatId': message.chat_id,
                'threadId': message.thread_id,
            })

        await run.on_async('inference', on_inference)

        await self.run(run)
        return run

    def _role_for_message(self, message):
        sender = message.sender
        is_digit = isinstance(sender, (int, str)) and str(sender).isdigit()
        if is_digit and str(sender) == str(self.id):
            return 'assistant'
        elif sender == 'user':
            return 'user'
        elif sender.startswith("tmp-"):
            return 'assistant'
        else:
            return 'user'
    
    # Returns all resource values matching the kind, and the key when given
    def get_resources(self, kind, key=None):
        return [r for r in self.resources
                if r.kind == kind and (key is None or r.key == key)]

    def get_tools(self, run=None):
        remotes = list(self.remotes)
        if run is not None and 'remotes' in run.input:
            remotes += list(run.input['remotes'])

        tools = list(self.tools) + self.app.getRemoteTools(self.user_id, remotes)

        # Dedupe by tool name (agent-configured vs run-input vs native collisions)
        seen, result = set(), []
        for tool in tools:
            if tool.name in seen:
                continue
            seen.add(tool.name)
            result.append(tool)

        return result


class DebugAgent(Agent):
    identifier = "debug"
    schema_fields = []
    tools = []

    async def received_message(self, message):
        chat_view = await self.app.get_chat_view(message.chat_id)
        response = chat_view.debug_string()

        await self.app.send_chat_message(
            chat_id=message.chat_id,
            sender=self.id,
            content=response,
            thread_id=chat_view.current_thread_id,
        )

# Holds parameters, outputs, and details of an agent call
class AgentRun:
    @classmethod
    def create(cls, app, agent_id, user_id, input=None):
        self = cls()
        self.app = app
        self.id = generate_tmp_id(app)
        self.agent_id = agent_id
        self.user_id = user_id
        self.done = False
        self.start_time = int(time.time())
        self.end_time = None

        # I/O envelope (no generic data bag)
        self.input = input or {}
        self.events = []
        self.event_listeners = {}

        # Default Streams
        self.output = None
        self.reasoning = None
        
        # Cancellation tracking
        # You should use 'cancelled', task and thread are for
        # external tasks that cannot be cancelled reliably
        self.task = None  # For asyncio tasks
        self.thread = None
        self.cancelled = False

        app.tmp.runs[self.id] = self
        
        # Emit run start event
        self.emit_event('created', {
            "agentId": self.agent_id,
            "userId": self.user_id,
            "startTime": self.start_time,
            "input": self.input,
        })
        
        return self

    def on(self, event_name, callback):
        """Register a callback for a named event.

        Replays every past event of this type to the callback immediately
        (with only the data payload), so late listeners catch up on history.
        """
        if event_name not in self.event_listeners:
            self.event_listeners[event_name] = []
        self.event_listeners[event_name].append(callback)

        # Replay: feed only the data payload of past events
        for event in self.events:
            if event["name"] == event_name:
                callback(event["data"])
    
    async def on_async(self, event_name, callback):
        event_loop = asyncio.get_event_loop()
        def new_callback(data):
            event_loop.create_task(callback(data))
        self.on(event_name, new_callback)

    def emit_event(self, event_name, data=None):
        """Record an event and notify existing listeners.

        Listeners receive only the *data* payload, not the full envelope.
        """
        event = {"name": event_name, "data": data, "time": time.time()}
        self.events.append(event)
        listeners = self.event_listeners.get(event_name, [])
        for cb in listeners:
            try:
                cb(data)
            except Exception as e:
                print(f"ERROR: handling event '{event_name}' in run {self.id}\n" + str(e))
                print(traceback.format_exc())

        # Broadcast event to websocket listeners
        self.app.emit_event("run-event", {
            "runId": self.id,
            "name": event_name,
            "data": data
        })
    
    def add_tools(self, tools):
        if "tools" in self.input:
            self.input["tools"] += tools
        else:
            self.input["tools"] = tools

    def add_tool_call(self, tool_name, arguments):
        self.emit_event("tool_call", {"name": tool_name, "arguments": arguments})

    def _serialize_input(self, input):
        """Convert ToolSpec objects to serializable dictionaries in input data."""
        if not input:
            return input

        result = {}
        for key, value in input.items():
            if key == "tools" and isinstance(value, list):
                # Convert ToolSpec objects to their OpenAI dictionary format
                result[key] = [tool.to_openai_dict() if hasattr(tool, 'to_openai_dict')
                    else tool for tool in value]
            else:
                result[key] = value
        return result

    def serialize(self):
        """Serialize the run object to a dictionary for JSON responses."""
        return {
            "id": self.id,
            "agentId": self.agent_id,
            "userId": self.user_id,
            "done": self.done,
            "startTime": self.start_time,
            "endTime": self.end_time,
            "input": self._serialize_input(self.input),
            "events": self.events,
            "output": self.output,
        }

    def cancel(self):
        """Cancel the agent run and stop any ongoing inference."""
        if self.cancelled or self.done:
            return
        
        self.cancelled = True
        print(f"Cancelling agent run {self.id}")
        
        # Cancel asyncio task if exists
        if self.task and not self.task.done():
            self.task.cancel()
        
        # Thread cannot be forcefully cancelled, but we can signal
        if self.thread and self.thread.is_alive():
            # The thread should be checking self.cancelled in its loop
            print(f"Signalled thread to stop for run {self.id}")
        
        # Close the stream and mark as done
        self.close()
        
        # Emit cancellation event
        self.emit_event("cancelled", {})
        
        # Broadcast cancellation event
        self.app.emit_event("run-cancelled", {
            "runId": self.id,
            "agentId": self.agent_id,
            "endTime": int(time.time()),
        })

    def close(self):
        if self.done:
            return
        self.done = True
        self.end_time = int(time.time())
        if self.output: self.output.close()
        if self.reasoning: self.reasoning.close()
        self.emit_event("done", {})