# Copyright 2019 The Feast Authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import concurrent.futures
import contextlib
import contextvars
import dataclasses
import hashlib
import logging
import os
import platform
import sys
import typing
import uuid
from datetime import datetime
from functools import wraps
from os.path import expanduser, join
from pathlib import Path

import requests

from feast import flags_helper
from feast.constants import DEFAULT_FEAST_USAGE_VALUE, FEAST_USAGE
from feast.version import get_version

USAGE_ENDPOINT = "https://usage.feast.dev"

_logger = logging.getLogger(__name__)
_executor = concurrent.futures.ThreadPoolExecutor(max_workers=3)

_is_enabled = os.getenv(FEAST_USAGE, default=DEFAULT_FEAST_USAGE_VALUE) == "True"

_constant_attributes = {
    "project_id": "",
    "session_id": str(uuid.uuid4()),
    "installation_id": None,
    "version": get_version(),
    "python_version": platform.python_version(),
    "platform": platform.platform(),
    "env_signature": hashlib.md5(
        ",".join(
            sorted([k for k in os.environ.keys() if not k.startswith("FEAST")])
        ).encode()
    ).hexdigest(),
}

APPLICATION_NAME = "feast-dev/feast"
USER_AGENT = "{}/{}".format(APPLICATION_NAME, get_version())


def get_user_agent():
    return USER_AGENT


def set_current_project_uuid(project_uuid: str):
    _constant_attributes["project_id"] = project_uuid


@dataclasses.dataclass
class FnCall:
    fn_name: str
    id: str

    start: datetime
    end: typing.Optional[datetime] = None

    parent_id: typing.Optional[str] = None


class Sampler:
    def should_record(self) -> bool:
        raise NotImplementedError

    @property
    def priority(self):
        return 0


class AlwaysSampler(Sampler):
    def should_record(self) -> bool:
        return True


class RatioSampler(Sampler):
    MAX_COUNTER = (1 << 32) - 1

    def __init__(self, ratio):
        assert 0 < ratio <= 1, "Ratio must be within (0, 1]"
        self.ratio = ratio
        self.total_counter = 0
        self.sampled_counter = 0

    def should_record(self) -> bool:
        self.total_counter += 1
        if self.total_counter == self.MAX_COUNTER:
            self.total_counter = 1
            self.sampled_counter = 1

        decision = self.sampled_counter < self.ratio * self.total_counter
        self.sampled_counter += int(decision)
        return decision

    @property
    def priority(self):
        return int(1 / self.ratio)


class UsageContext:
    attributes: typing.Dict[str, typing.Any]

    call_stack: typing.List[FnCall]
    completed_calls: typing.List[FnCall]

    exception: typing.Optional[Exception] = None
    traceback: typing.Optional[typing.Tuple[str, int, str]] = None

    sampler: Sampler = AlwaysSampler()

    def __init__(self):
        self.attributes = {}
        self.call_stack = []
        self.completed_calls = []


_context = contextvars.ContextVar("usage_context", default=UsageContext())


def _set_installation_id():
    if os.getenv("FEAST_FORCE_USAGE_UUID"):
        _constant_attributes["installation_id"] = os.getenv("FEAST_FORCE_USAGE_UUID")
        _constant_attributes["installation_ts"] = datetime.utcnow().isoformat()
        return

    feast_home_dir = join(expanduser("~"), ".feast")
    installation_timestamp = datetime.utcnow()

    try:
        Path(feast_home_dir).mkdir(exist_ok=True)
        usage_filepath = join(feast_home_dir, "usage")

        if os.path.exists(usage_filepath):
            installation_timestamp = datetime.utcfromtimestamp(
                os.path.getmtime(usage_filepath)
            )
            with open(usage_filepath, "r") as f:
                installation_id = f.read()
        else:
            installation_id = str(uuid.uuid4())

            with open(usage_filepath, "w") as f:
                f.write(installation_id)
            print(
                "Feast is an open source project that collects "
                "anonymized error reporting and usage statistics. To opt out or learn"
                " more see https://docs.feast.dev/reference/usage"
            )
    except OSError as e:
        _logger.debug(f"Unable to configure usage {e}")
        installation_id = "undefined"

    _constant_attributes["installation_id"] = installation_id
    _constant_attributes["installation_ts"] = installation_timestamp.isoformat()


_set_installation_id()


def _export(event: typing.Dict[str, typing.Any]):
    _executor.submit(requests.post, USAGE_ENDPOINT, json=event, timeout=2)


def _produce_event(ctx: UsageContext):
    if ctx.sampler and not ctx.sampler.should_record():
        return
    # Cannot check for unittest because typeguard pulls in unittest
    is_test = flags_helper.is_test() or bool({"pytest"} & sys.modules.keys())
    event = {
        "timestamp": datetime.utcnow().isoformat(),
        "is_test": is_test,
        "is_webserver": (
            not is_test and bool({"uwsgi", "gunicorn", "fastapi"} & sys.modules.keys())
        ),
        "calls": [
            dict(
                fn_name=c.fn_name,
                id=c.id,
                parent_id=c.parent_id,
                start=c.start and c.start.isoformat(),
                end=c.end and c.end.isoformat(),
            )
            for c in reversed(ctx.completed_calls)
        ],
        "entrypoint": ctx.completed_calls[-1].fn_name,
        "exception": repr(ctx.exception) if ctx.exception else None,
        "traceback": ctx.traceback if ctx.exception else None,
        **_constant_attributes,
    }
    event.update(ctx.attributes)
    _export(event)


@contextlib.contextmanager
def tracing_span(name):
    """
    Context manager for wrapping heavy parts of code in tracing span
    """
    if _is_enabled:
        ctx = _context.get()
        if not ctx.call_stack:
            raise RuntimeError("tracing_span must be called in usage context")

        last_call = ctx.call_stack[-1]
        fn_call = FnCall(
            id=uuid.uuid4().hex,
            parent_id=last_call.id,
            fn_name=f"{last_call.fn_name}.{name}",
            start=datetime.utcnow(),
        )
    try:
        yield
    finally:
        if _is_enabled:
            fn_call.end = datetime.utcnow()
            ctx.completed_calls.append(fn_call)


def log_exceptions_and_usage(*args, **attrs):
    """
    This function decorator enables three components:
    1. Error tracking
    2. Usage statistic collection
    3. Time profiling

    This data is being collected, anonymized and sent to Feast Developers.
    All events from nested decorated functions are being grouped into single event
    to build comprehensive context useful for profiling and error tracking.

    Usage example (will result in one output event):
        @log_exceptions_and_usage
        def fn(...):
            nested()

        @log_exceptions_and_usage(attr='value')
        def nested(...):
            deeply_nested()

        @log_exceptions_and_usage(attr2='value2', sample=RateSampler(rate=0.1))
        def deeply_nested(...):
            ...
    """
    sampler = attrs.pop("sampler", AlwaysSampler())

    def clear_context(ctx):
        _context.set(UsageContext())  # reset context to default values
        # TODO: Figure out why without this, new contexts.get aren't reset
        ctx.call_stack = []
        ctx.completed_calls = []
        ctx.attributes = {}

    def decorator(func):
        if not _is_enabled:
            return func

        @wraps(func)
        def wrapper(*args, **kwargs):
            ctx = _context.get()
            ctx.call_stack.append(
                FnCall(
                    id=uuid.uuid4().hex,
                    parent_id=ctx.call_stack[-1].id if ctx.call_stack else None,
                    fn_name=_fn_fullname(func),
                    start=datetime.utcnow(),
                )
            )
            ctx.attributes.update(attrs)

            try:
                return func(*args, **kwargs)
            except Exception:
                if ctx.exception:
                    # exception was already recorded
                    raise

                _, exc, traceback = sys.exc_info()
                ctx.exception = exc
                ctx.traceback = _trace_to_log(traceback)

                if traceback:
                    raise exc.with_traceback(traceback)

                raise exc
            finally:
                ctx.sampler = (
                    sampler if sampler.priority > ctx.sampler.priority else ctx.sampler
                )
                last_call = ctx.call_stack.pop(-1)
                last_call.end = datetime.utcnow()
                ctx.completed_calls.append(last_call)

                if not ctx.call_stack or (
                    len(ctx.call_stack) == 1
                    and "feast.feature_store.FeatureStore.serve"
                    in str(ctx.call_stack[0].fn_name)
                ):
                    # When running `feast serve`, the serve method never exits so it gets
                    # stuck otherwise
                    _produce_event(ctx)
                    clear_context(ctx)

        return wrapper

    if args:
        return decorator(args[0])

    return decorator


def log_exceptions(*args, **attrs):
    """
    Function decorator that track errors and send them to Feast Developers
    """

    def decorator(func):
        if not _is_enabled:
            return func

        @wraps(func)
        def wrapper(*args, **kwargs):
            if _context.get().call_stack:
                # we're already inside usage context
                # let it handle exception
                return func(*args, **kwargs)

            fn_call = FnCall(
                id=uuid.uuid4().hex, fn_name=_fn_fullname(func), start=datetime.utcnow()
            )
            try:
                return func(*args, **kwargs)
            except Exception:
                _, exc, traceback = sys.exc_info()

                fn_call.end = datetime.utcnow()

                ctx = UsageContext()
                ctx.exception = exc
                ctx.traceback = _trace_to_log(traceback)
                ctx.attributes = attrs
                ctx.completed_calls.append(fn_call)
                _produce_event(ctx)

                if traceback:
                    raise exc.with_traceback(traceback)

                raise exc

        return wrapper

    if args:
        return decorator(args[0])

    return decorator


def set_usage_attribute(name, value):
    """
    Extend current context with custom attribute
    """
    ctx = _context.get()
    ctx.attributes[name] = value


def _trim_filename(filename: str) -> str:
    return filename.split("/")[-1]


def _fn_fullname(fn: typing.Callable):
    return fn.__module__ + "." + fn.__qualname__


def _trace_to_log(traceback):
    log = []
    while traceback is not None:
        log.append(
            (
                _trim_filename(traceback.tb_frame.f_code.co_filename),
                traceback.tb_lineno,
                traceback.tb_frame.f_code.co_name,
            )
        )
        traceback = traceback.tb_next

    return log
