import asyncio
from datetime import datetime

import database as db 
import temporary as tmp
import agent as agent_mod
import agents as agents_mod

class AgentMixin:
    # --- Base ---

    async def get_agent(self, agent_id):
        if tmp.is_tmp_id(agent_id):
            agent = self.tmp.agents.get(agent_id)
        else:
            agent = await db.Agent.get_or_none(id=agent_id)

        if not agent:
            raise Exception(f"Agent not found with id {agent_id}")
        return agent

    async def agent_id_from_name(self, agentname: str):
        agent_id = await db.Agent.get_or_none(agentname=agentname).values_list("id", flat=True)
        if not agent_id:
            raise Exception(f"Agent not found with agentname {agentname}")
        return agent_id

    async def create_agent(self, user_id, name, agentname, description, picture, kind, data, main_chat_id=None, temporary=False):
        print("create_agent; user_id, name:", user_id, name)

        if temporary:
            agent = tmp.Agent.create(
                self,
                user_id=user_id,
                name=name,
                agentname=agentname,
                main_chat_id=main_chat_id,
                description=description,
                picture=picture,
                kind=kind,
                data=data or {}
            )

            if main_chat_id is not None:
                if tmp.is_tmp_id(main_chat_id):
                    self.tmp.chats[main_chat_id].participants.append(agent.id)
                else:
                    print(f"ERROR: cannot use {main_chat_id} as main chat for temporary agent {agentname or agent.id}")
                    agent.main_chat_id = None

            return agent.id

        if main_chat_id is None:
            chat = await db.Chat.create(
                user_id=user_id,
                name=name
            )
            main_chat_id = chat.id

        agent = await db.Agent.create(
            user_id=user_id,
            name=name,
            agentname=agentname,
            main_chat_id=main_chat_id,
            description=description,
            picture=picture,
            kind=kind,
            data=data or {}
        )
        print(f"Created agent #{agent.id}: {agent.name}")

        print(f"... adding agent {repr(agent.id)} to chat {repr(main_chat_id)}")
        await db.Membership.create(
            chat_id=main_chat_id,
            agent_id=agent.id
        )
        return agent.id

    async def update_agent(self, agent_id, name=None, description=None, picture=None, agentname=None, data=None):
        agent = await self.get_agent(agent_id)

        summary = []
        if name is not None:
            agent.name = name
            summary.append('name')
        if description is not None:
            agent.description = description
            summary.append('description')
        if picture is not None:
            agent.picture = picture
            summary.append('picture')
        if agentname is not None:
            agent.agentname = agentname
            summary.append('agentname')

        if data:
            # Tortoise tracks changes in JSONField automatically
            agent.data.update(data)
            summary.extend([f"data.{k}" for k in data.keys()])

        if db.is_instance(agent):
            await agent.save()
        print(f"Updated agent #{agent.id} fields: {summary}")

    async def get_agents(self, user_id):
        agents = await db.Agent.filter(user_id=user_id).all()
        if user_id in self.tmp.user_agents:
            agents += [self.tmp.agents[agent_id] for agent_id in self.tmp.user_agents[user_id]]

        return [{
            "id": agent.id,
            "name": agent.name,
            "agentname": agent.agentname,
            "description": agent.description,
            "picture": agent.picture,
            "kind": agent.kind,
        } for agent in agents]

    async def get_agent_resources(self, agent_id, kind=None, key=None):
        filter = {"agent_id": agent_id}
        if kind: filter['kind'] = kind
        if key: filter['key'] = key

        values = await db.AgentResource.filter(**filter)
        return values

    async def save_agent_resource(self, agent_id, kind, key, value):
        await db.AgentResource.update_or_create(
            {'value': value},
            agent_id=agent_id,
            kind=kind,
            key=key,
        )
        print(f"Saved agent resource: {agent_id} - {kind} - {key}")

    async def delete_agent_resource(self, agent_id, kind, key):
        await db.AgentResource.filter(
            agent_id=agent_id,
            kind=kind,
            key=key,
        ).delete()
        print(f"Deleted agent resource: {agent_id} - {kind} - {key}")

    # --- Control & Actions ---

    async def get_agent_control(self, agent_id):
        agent = await self.get_agent(agent_id)
        resources = await self.get_agent_resources(agent_id)

        agent_class = next((kind for kind in self.agent_classes if kind.identifier == agent.kind), None)
        if agent_class is None:
            raise Exception(f"invalid agent kind: {agent.kind}")

        control = agent_class(self, agent.data, resources)
        control._set_basic_info(agent_id, agent.name, agent.user_id)
        return control

    async def create_agent_run(self, agent_id):
        agent = await self.get_agent(agent_id)
        run = agent_mod.AgentRun.create(
            self,
            agent_id=agent_id,
            user_id=agent.user_id,
        )
        return run.id

    async def cancel_agent_run(self, run_id):
        """Cancel an ongoing agent run by ID."""
        run = self.tmp.runs.get(run_id)
        if run:
            run.cancel()
            return True
        return False

    async def get_agent_runs(self, agent_id=None, user_id=None):
        runs = list(self.tmp.runs.values())
        if agent_id:
            runs = [r for r in runs if r.agent_id == agent_id]
        elif user_id:
            runs = [r for r in runs if r.user_id == user_id]
        return [r.serialize() for r in runs]

    async def prompt_agent(self, agent_id, instruction, input=None, options={}):
        # Mirror agent.prompt's generator so run ownership and lifecycle stay
        # inside the mixin: callers just iterate and receive the same raw
        # inference.Message objects. get_agent_control is awaited lazily on first
        # iteration, keeping the control object alive for the generator's
        # lifetime.

        def normalize_tool(t):
            # this function makes a lot of sense in ToolSpec itself...
            from tool import ToolSpec, ToolParameter
            if isinstance(t, dict):
                params = []

                d = t.get('parameters', {})
                required = d.get('required', [])
                props = d.get('properties', {})

                for key, p in props.items():
                    params.append(ToolParameter(
                        name=key,
                        type=p.get('type'),
                        description=p.get('description'),
                        required=key in required,
                    ))
                return ToolSpec(
                    name=t['name'],
                    description=t['description'],
                    parameters=params,
                )
            else:
                return t

        if 'tools' in options:
            options = {**options}
            options['tools'] = [
                normalize_tool(t) for t in options['tools']
            ]

        control = await self.get_agent_control(agent_id)
        async for msg in control.prompt(instruction, input, **options):
            yield msg

    async def call_agent_action(self, agent_id, action, args):
        control = await self.get_agent_control(agent_id)
        response = await control.execute_action(action, args)
        return response

    # --- Information ---

    def get_agent_kinds(self):
        return [kind.identifier for kind in self.agent_classes]

    def get_agent_kind_info(self, kind):
        agent_class = next((it for it in self.agent_classes if it.identifier == kind), None)
        if agent_class is None:
            raise Exception(f"invalid agent kind: {kind}")
        return {
            "name": agent_class.identifier,
            "description": getattr(agent_class, 'description', None),
        }

    async def get_agent_config_fields(self, agent_id):
        agent = await self.get_agent(agent_id)

        agent_class = next((it for it in self.agent_classes if it.identifier == agent.kind), None)
        if not agent_class:
            raise Exception(f"Invalid agent kind: {agent.kind}")

        data = agent.data or {}
        return [{
            'name': field.name,
            'type': field.type,
            'value': data.get(field.name, field.default)
        } for field in agent_class.schema_fields]

    async def generate_default_user_data(self, user_id):
        chat_id = await self.create_chat(
            user_id=user_id,
            name="Initial Chat",
            description="First user chat with Intro.",
            picture=None,
            members=[],
        )
        await self.create_agent(
            user_id=user_id,
            name="Intro",
            agentname="intro",
            main_chat_id=chat_id,
            description="Introduces you to the application.",
            picture=None,
            kind="conversation",
            data={}
        )