

from arkitekt.agents.base import Agent
from arkitekt.messages.postman.provide.provide_log import ProvideLogMessage
from arkitekt.messages.postman.log import LogLevel
from arkitekt.monitor.monitor import AgentPanel, get_current_monitor
from herre.console.context import get_current_console
from arkitekt.packers.utils import shrink_outputs
from re import template
from arkitekt.messages.postman.assign.assign_cancelled import AssignCancelledMessage
from arkitekt.messages.postman.unassign.bounced_forwarded_unassign import BouncedForwardedUnassignMessage
from arkitekt.messages.postman.assign.assign_critical import AssignCriticalMessage
from typing import Dict
from arkitekt.messages.postman.assign.assign_return import AssignReturnMessage
from arkitekt.messages.postman.assign.bounced_forwarded_assign import BouncedForwardedAssignMessage
from arkitekt.messages.postman.assign.bounced_assign import BouncedAssignMessage
from arkitekt.messages.base import MessageModel
from arkitekt.messages.postman.provide.provide_transition import ProvideState
from arkitekt.transport.base import Transport
from arkitekt.packers.transpilers.base import Transpiler
from arkitekt.messages.postman.provide.bounced_provide import BouncedProvideMessage
import asyncio
from asyncio.tasks import Task
from arkitekt.messages.postman.provide import ProvideTransitionMessage
import logging
from arkitekt.schema.template import Template
from koil.koil import Koil, get_current_koil
from koil.loop import koil



logger = logging.getLogger(__name__)



class Actor:
    template: Template

    def __init__(self, *args, koil: Koil = None, strict=False, expand_inputs=True, shrink_outputs=True, transpilers= {}, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.koil = koil or get_current_koil()
        self.loop = self.koil.loop
        self.strict = strict
        self.expand_inputs = expand_inputs
        self.shrink_outputs = shrink_outputs
        self.runningAssignments: Dict[str, asyncio.Task] = {} # Running assignments indexed by assignment reference


    async def log(self, message, level = LogLevel.INFO):
        logger.info("{level}: {message}")


    async def provide_log(self, message: str, level = LogLevel.INFO):
        await self.agent.transport.forward(ProvideLogMessage(
                data= {
                    "level": level,
                    "message": message
                },
                meta=self.provision.meta.dict(exclude={"type"})
            ))
        logger.info(f"Provide Log: {level}: {message}")


    def run(self, *args, **kwargs):
        return koil(self.arun(*args, **kwargs))

    async def acall(self, message: MessageModel):
        await self.in_queue.put(message)


    async def on_provide(self, message: BouncedProvideMessage):
        print("Getting Called Here")
        return None


    async def on_unprovide(self, message: BouncedProvideMessage):
        return None


    async def on_assign(self, message: BouncedForwardedAssignMessage):
        raise NotImplementedError("Needs to be owerwritten in Actor Subclass")


    async def arun(self, provision: BouncedProvideMessage, agent: Agent):
        self.provision = provision
        self.agent = agent
        self.transport = agent.transport
        self.in_queue = asyncio.Queue()

        try:
            self.template = await Template.asyncs.get(id=self.provision.data.template)
            

            await self.transport.forward(ProvideTransitionMessage(
                data= {
                    "state": ProvideState.PROVIDING,
                    "message": "We just got started Bay"
                },
                meta={
                    "reference": self.provision.meta.reference,
                    "extensions": self.provision.meta.extensions
                }
            ))



            await self.on_provide(self.provision)

            await self.transport.forward(ProvideTransitionMessage(
                data= {
                    "state": ProvideState.ACTIVE,
                    "message": "We just got started Bay"
                },
                meta={
                    "reference": self.provision.meta.reference,
                    "extensions": self.provision.meta.extensions
                }
            ))

            while True:
                await self.log("Waiting for assignmements")
                message = await self.in_queue.get()
                logger.info(f"Received Message {message}")
                
                if isinstance(message, BouncedForwardedAssignMessage):
                    await self.log("Assigningment received")
                    task = self.loop.create_task(self.on_assign(message))
                    self.runningAssignments[message.meta.reference] = task

                if isinstance(message, BouncedForwardedUnassignMessage):
                    if message.data.assignation in self.runningAssignments:
                        task = self.runningAssignments[message.data.assignation]
                        if not task.done():
                            logger.info("Task is being cancelled")
                            task.cancel()
                        else:
                            logger.error("Task was already done")
                    else:
                        logger.error("Task was never assigned to this actor")
                        if self.strict: raise Exception("Received cancellation for Task that was never assinged to this actor!")
                        await self.transport.forward(AssignCancelledMessage(data={
                            "canceller": str("Cancelled because actor receiving this cancellation never had this task but was also not strict")
                        }, meta = {
                            "reference": message.data.assignation,
                            "extensions": message.meta.extensions
                        }))



        except Exception as e:
            logger.exception(e)
            await self.log(f"Provision Exception {str(e)}")
            await self.transport.forward(ProvideTransitionMessage(
                data= {
                    "state": ProvideState.CRITICAL,
                    "message": f"{e}"
                },
                meta={
                    "reference": self.provision.meta.reference,
                    "extensions": self.provision.meta.extensions
                }
            ))


        except asyncio.CancelledError as e:
            await self.transport.forward(ProvideTransitionMessage(
                data= {
                    "state": ProvideState.CANCELING,
                    "message": f"{e}"
                },
                meta={
                    "reference": self.provision.meta.reference,
                    "extensions": self.provision.meta.extensions
                }
            ))

            await self.on_unprovide(self.provision)


            logger.info("Doing Whatever needs to be done to cancel!")
            await self.transport.forward(ProvideTransitionMessage(
                data= {
                    "state": ProvideState.CANCELLED,
                    "message": f"{e}"
                },
                meta={
                    "reference": self.provision.meta.reference,
                    "extensions": self.provision.meta.extensions
                }
            ))
            raise e





