import asyncio
from collections import defaultdict
import contextvars
import copy
from functools import wraps
import logging
import time
from typing import Any, Callable, DefaultDict, List, TypeVar, cast
import uuid

from lytix_py.MetricCollector.MetricCollector import MetricCollector

"""
@see https://mypy.readthedocs.io/en/stable/generics.html#declaring-decorators
"""
F = TypeVar("F", bound=Callable[..., Any])


_lytixTestContext: contextvars.ContextVar[DefaultDict[str, str]] = (
    contextvars.ContextVar(
        "lytixContext",
        default=defaultdict(
            lambda: {
                "traceId": None,
                "modelName": None,
                "modelInput": None,
                "modelOutput": None,
            },
        ),
    )
)

_lytixTraceContext: contextvars.ContextVar[DefaultDict[str, str]] = (
    contextvars.ContextVar(
        "lytixTraceContext",
        default=defaultdict(
            lambda: {
                "traceId": None,
                "modelName": None,
                "modelInput": None,
                "modelOutput": None,
                "userIdentifier": None,
                "sessionId": None,
            },
        ),
    )
)


class LytixWrapper:
    _logger = logging.getLogger("lytix")
    _testState = {}

    def trace(
        self,
        *kwargs,
        modelName: str,
    ) -> Callable[[F], F]:
        """
        Main trace decorator.
        @param modelName: Model name we are tracing
        """

        def decorator(func: F) -> F:
            """
            First lets see if this is an async function
            """
            isAsyncFunction = asyncio.iscoroutinefunction(func)

            """
            Wrap our function in a context state
            """
            if isAsyncFunction:
                return self._asyncTraceWrapper(func, modelName)
            else:
                return self._syncTraceWrapper(func, modelName)

        return decorator

    def test(
        self,
        *kwargs,
        testsToRun: List[str],
        _testSuiteId: str = None,
        _testName: str = None,
    ) -> Callable[[F], F]:
        """
        Main testing decorator.
        @param testGroupId: This is set when the test suite runs
        @param testName: This is set when the test suite runs
        """

        def decorator(func: F) -> F:
            """
            If we have any function args, it can't be a test
            """
            if kwargs:
                self._logger.error("Lytix tests cannot have function args")
            else:
                self._testState[func.__name__] = {
                    "function": func,
                    "testArgs": {"testsToRun": testsToRun},
                }

            """
            First lets see if this is an async function
            """
            isAsyncFunction = asyncio.iscoroutinefunction(func)

            """
            Wrap our function in a context state
            """
            if isAsyncFunction:
                return self._asyncTestWrapper(func, testsToRun, _testSuiteId, _testName)
            else:
                return self._syncTestWrapper(func, testsToRun, _testSuiteId, _testName)

        return decorator

    def _syncTraceWrapper(self, func: F, modelName: str):
        """
        Sync wrapper for the trace decorator
        """

        @wraps(func)
        def syncDecorator(*args, **kwargs):
            traceContext = self._prepTraceDecoratorContext(modelName)
            startTime = time.time()
            result = func(*args, **kwargs)
            responseTime = int(
                (time.time() - startTime) * 1000
            )  # Convert to milliseconds

            """
            If we have model IO set, send it to HQ
            """
            self.reportModelIO(responseTime, traceContext)

            return result

        return cast(F, syncDecorator)

    def _asyncTraceWrapper(self, func: F, modelName: str):
        """
        Async wrapper for the trace decorator
        """

        @wraps(func)
        async def asyncDecorator(*args, **kwargs):
            traceContext = self._prepTraceDecoratorContext(modelName)
            startTime = time.time()
            result = await func(*args, **kwargs)
            responseTime = int(
                (time.time() - startTime) * 1000
            )  # Convert to milliseconds

            """
            If we have model IO set, send it to HQ
            """
            self.reportModelIO(responseTime, traceContext)

            return result

        return cast(F, asyncDecorator)

    def reportModelIO(self, responseTime: int, traceContexT: dict):
        """
        Reports model input and output to HQ
        """

        context, valid = self._getTraceContext()
        if not valid:
            return

        modelInput = context["modelInput"]
        modelOutput = context["modelOutput"]
        modelName = context["modelName"]

        if modelInput is None or modelOutput is None or modelName is None:
            self._logger.error(
                "Missing modelInput, modelOutput, or modelName in context, make sure you have setup the correct decorator"
            )
            return

        """
        Pull optional fields now as well
        """
        userIdentifier = context["userIdentifier"]
        sessionId = context["sessionId"]

        """
        Send to HQ
        """
        MetricCollector.captureModelIO(
            modelName=modelName,
            modelInput=copy.deepcopy(modelInput),
            modelOutput=copy.deepcopy(modelOutput),
            userIdentifier=userIdentifier,
            sessionId=sessionId,
            modelResponseTime=responseTime,
        )

    def _syncTestWrapper(
        self, func: F, testsToRun: List[str], testSuiteId: str, testName: str
    ):
        """
        Sync wrapper for the decorator
        """

        @wraps(func)
        def syncDecorator(*args, **kwargs):
            testContext = self._prepTestDecoratorContext(testsToRun)
            func(*args, **kwargs)

            """
            Before returning to the user, send the test + IO back to HQ
            """
            testResult = self._startAndWaitForLytixTest(testSuiteId, testName)

            return testResult

        return cast(F, syncDecorator)

    def _asyncTestWrapper(
        self, func: F, testsToRun: List[str], testSuiteId: str, testName: str
    ):
        """
        Async wrapper for the decorator
        """

        @wraps(func)
        async def asyncDecorator(*args, **kwargs):
            testContext = self._prepTestDecoratorContext(testsToRun)
            await func(*args, **kwargs)

            """
            Before returning to the user, send the test + IO back to HQ
            """
            testResult = self._startAndWaitForLytixTest(testSuiteId, testName)

            return testResult

        return cast(F, asyncDecorator)

    def _startAndWaitForLytixTest(self, testSuiteId: str, testName: str):
        """
        Starts a new test in Lytix and waits for it to complete
        """
        context, valid = self._getTestContext()
        if not valid:
            return

        """
        Validate we have everything we need
        """
        input = context["modelInput"]
        output = context["modelOutput"]
        testsToRun = context["testsToRun"]
        traceId = context["traceId"]
        if input is None or output is None or testsToRun is None or traceId is None:
            self._logger.error(
                "Missing input, output, testsToRun, or traceId in context, make sure you have setup the correct decorator"
            )
            return False

        """
        Start the test
        """
        testRunIds = MetricCollector._kickoffTestRun(
            input, output, testsToRun, traceId, testSuiteId, testName
        )

        if testRunIds is False:
            return "FAILED"

        """
        Wait for it to finish
        """
        allFinished = False
        anyFailed = False
        while allFinished is False:
            allFinished = True
            """
            Loop over all testRunIds
            """
            for testRunId in testRunIds:
                testStatus = MetricCollector._getTestStatus(testRunId)
                if testStatus["testResult"] == "FAILED":
                    anyFailed = True
                if testStatus["testStatus"] not in ["FAILED", "SUCCESS"]:
                    allFinished = False
                    break

            if allFinished is False:
                time.sleep(5)

        """
        Result can be FAILED or SUCCESS, if it's FAILED
        """
        if anyFailed:
            return "FAILED"
        else:
            return "SUCCESS"

    def _prepTestDecoratorContext(self, testsToRun: List[str]) -> dict:
        """
        Prepares the decorator context
        """
        try:
            stack = _lytixTestContext.get().copy()
            stack["testsToRun"] = testsToRun
            stack["traceId"] = str(uuid.uuid4())
            stack["modelOutput"] = None
            stack["modelInput"] = None
            stack["modelName"] = None
            _lytixTestContext.set(stack)

            return {"traceId": stack["traceId"]}
        except Exception as e:
            self._logger.error(f"Error setting up decorator context {e}")

    def _prepTraceDecoratorContext(self, modelName: str) -> dict:
        """
        Prepares the decorator context
        """
        try:
            stack = _lytixTraceContext.get().copy()
            stack["modelName"] = modelName
            stack["traceId"] = str(uuid.uuid4())
            stack["modelOutput"] = None
            stack["modelInput"] = None
            stack["userIdentifier"] = None
            stack["sessionId"] = None
            _lytixTraceContext.set(stack)

            return {"traceId": stack["traceId"]}
        except Exception as e:
            self._logger.error(f"Error setting up trace decorator context {e}")

    def _getTestContext(self):
        """
        Gets the current context
        """
        existingContext = _lytixTestContext.get()
        if existingContext["traceId"] is None:
            self._logger.debug(
                "No traceId in context, make sure you have setup the correct decorator"
            )
            return {}, False
        return existingContext, True

    def _getTraceContext(self):
        """
        Gets the current context
        """
        existingContext = _lytixTraceContext.get()
        if existingContext["traceId"] is None:
            self._logger.debug(
                "No traceId in context, make sure you have setup the correct decorator"
            )
            return {}, False
        return existingContext, True

    def setOutput(self, modelOutput: str):
        """
        Sets model output in the current context
        """
        # Try to set on Test context
        context, valid = self._getTestContext()
        if valid:
            context["modelOutput"] = modelOutput
            _lytixTestContext.set(context)

        #  Then set on trace context
        context, valid = self._getTraceContext()
        if valid:
            context["modelOutput"] = modelOutput
            _lytixTraceContext.set(context)

    def setInput(self, modelInput: str):
        """
        Sets model input in the current context
        """
        # Try to set on Test context
        context, valid = self._getTestContext()
        if valid:
            context["modelInput"] = modelInput
            _lytixTestContext.set(context)

        #  Then set on trace context
        context, valid = self._getTraceContext()
        if valid:
            context["modelInput"] = modelInput
            _lytixTraceContext.set(context)

    def setUserIdentifier(self, userIdentifier: str):
        """
        Sets user identifier in the current context
        """
        context, valid = self._getTraceContext()
        if valid:
            context["userIdentifier"] = userIdentifier
            _lytixTraceContext.set(context)

    def setSessionId(self, sessionId: str):
        """
        Sets session id in the current context
        """
        context, valid = self._getTraceContext()
        if valid:
            context["sessionId"] = sessionId
            _lytixTraceContext.set(context)


lytix = LytixWrapper()
