"""
Materialization service for dataset late materialization.
"""
import asyncio
import dataclasses
from datetime import UTC, datetime
import json
import logging
import time
from typing import Any, Self
import uuid

from prometheus_client import Info, start_http_server
import pymongo
import pymongo.errors
import motor.motor_asyncio
import requests.exceptions
from rest_tools.client import RestClient, ClientCredentialsAuth
from rest_tools.server import RestServer
from tornado.web import HTTPError
from tornado.web import RequestHandler as TornadoRequestHandler
from wipac_dev_tools import from_environment_as_dataclass

from iceprod.util import VERSION_STRING
from ..prom_utils import AsyncMonitor
from iceprod.rest.auth import authorization, attr_auth
from iceprod.rest.base_handler import IceProdRestConfig, APIBase
from iceprod.server.util import nowstr, datetime2str
from .service import MaterializationService

logger = logging.getLogger('server')


class BaseHandler(APIBase):
    def initialize(self, *args, materialization_service=None, rest_client=None, **kwargs):
        super().initialize(*args, **kwargs)
        self.materialization_service = materialization_service
        self.rest_client = rest_client

    async def new_request(self, args):
        # validate first
        fields = {
            'dataset_id': str,
            'set_status': str,
            'num': int,
        }
        if set(args)-set(fields):  # don't let random args through
            raise HTTPError(400, reason='invalid params')
        for k in fields:
            if k in args and not isinstance(args[k], fields[k]):
                r = 'key "{}" should be of type {}'.format(k, fields[k].__name__)
                raise HTTPError(400, reason=r)

        # set some fields
        now = nowstr()
        data = {
            'materialization_id': uuid.uuid1().hex,
            'status': 'waiting',
            'create_timestamp': now,
            'modify_timestamp': now,
            'creator': self.auth_data.get('username', None),
            'role': self.auth_data.get('role', None),
        }
        for k in fields:
            if k in args:
                data[k] = args[k]

        # insert
        await self.db.materialization.insert_one(data)
        return data

    async def check_attr_auth(self, arg, val, role):
        """
        Based on the request groups or username, check if they are allowed to
        access `arg`:`role`.

        Runs a remote query to the IceProd API.

        Args:
            arg (str): attribute name to check
            val (str): attribute value
            role (str): the role to check for (read|write)
        """
        assert self.rest_client
        args = {
            'name': arg,
            'value': val,
            'role': role,
            'username': self.current_user,
            'groups': self.auth_groups,
        }
        try:
            await self.rest_client.request('POST', '/auths', args)
        except requests.exceptions.HTTPError as e:
            if e.response.status_code == 403:
                raise HTTPError(403, 'auth failed')
            else:
                raise HTTPError(500, 'auth could not be completed')


class StatusHandler(BaseHandler):
    """
    Handle materialization status requests.
    """
    @authorization(roles=['admin', 'system'])
    async def get(self, materialization_id):
        """
        Get materialization status.

        If materialization_id is invalid, returns http code 404.

        Args:
            materialization_id (str): materialization request id

        Returns:
            dict: materialization metadata
        """
        ret = await self.db.materialization.find_one({'materialization_id':materialization_id},
                                                     projection={'_id':False})
        if not ret:
            self.send_error(404, reason="Materialization request not found")
        else:
            self.write(ret)


class RequestHandler(BaseHandler):
    """
    Handle basic materialization requests.
    """
    @authorization(roles=['admin', 'system'])
    async def post(self):
        """
        Create basic materialization request.

        On success, returns http code 201.

        Params:
            num: number of jobs to buffer per dataset

        Returns:
            dict: {result: materialization_id}
        """
        args = json.loads(self.request.body) if self.request.body else {}

        data = await self.new_request(args)

        # return success
        self.set_status(201)
        self.write({'result': data['materialization_id']})


class RequestDatasetHandler(BaseHandler):
    """
    Handle dataset materialization requests.
    """
    @authorization(roles=['admin', 'user', 'system'])
    @attr_auth(arg='dataset_id', role='write')
    async def post(self, dataset_id):
        """
        Create dataset materialization request.

        On success, returns http code 201.

        Params:
            num: number of jobs to buffer per dataset

        Args:
            dataset_id (str): dataset_id to materialize

        Returns:
            dict: {result: materialization_id}
        """
        args = json.loads(self.request.body) if self.request.body else {}
        args['dataset_id'] = dataset_id

        data = await self.new_request(args)

        # return success
        self.set_status(201)
        self.write({'result': data['materialization_id']})


class DatasetStatusHandler(BaseHandler):
    """
    Handle dataset materialization request statuses.
    """
    @authorization(roles=['admin', 'user', 'system'])
    @attr_auth(arg='dataset_id', role='read')
    async def get(self, dataset_id):
        """
        Get latest materialization status for a dataset, if any.

        If dataset_id is invalid, returns http code 404.

        Args:
            dataset_id (str): dataset id to check

        Returns:
            dict: materialization metadata
        """
        ret = await self.db.materialization.find_one(
            {'dataset_id': dataset_id},
            projection={'_id':False},
            sort=[('modify_timestamp', pymongo.DESCENDING)],
        )
        if not ret:
            self.send_error(404, reason="Materialization request not found")
        else:
            self.write(ret)


class HealthHandler(BaseHandler):
    """
    Handle health requests.
    """
    async def get(self):
        """
        Get health status.

        Returns based on exit code, 200 = ok, 400 = failure
        """
        assert self.materialization_service
        now = time.time()
        status = {
            'now': nowstr(),
            'start_time': datetime2str(datetime.fromtimestamp(self.materialization_service.start_time, tz=UTC)),
            'last_run_time': "",
            'last_success_time': "",
            'last_cleanup_time': "",
            'num_requests': -1,
        }
        try:
            if self.materialization_service.last_run_time is None and self.materialization_service.start_time + 3600 < now:
                self.send_error(500, reason='materialization was never run')
                return
            if self.materialization_service.last_run_time is not None:
                if self.materialization_service.last_run_time + 3600 < now:
                    self.send_error(500, reason='materialization has stopped running')
                    return
                status['last_run_time'] = datetime2str(datetime.fromtimestamp(self.materialization_service.last_run_time, tz=UTC))
            if self.materialization_service.last_success_time is None and self.materialization_service.start_time + 86400 < now:
                self.send_error(500, reason='materialization was never successful')
                return
            if self.materialization_service.last_success_time is not None:
                if self.materialization_service.last_success_time + 86400 < now:
                    self.send_error(500, reason='materialization has stopped being successful')
                    return
                status['last_success_time'] = datetime2str(datetime.fromtimestamp(self.materialization_service.last_success_time, tz=UTC))
            if self.materialization_service.last_cleanup_time is not None:
                status['last_cleanup_time'] = datetime2str(datetime.fromtimestamp(self.materialization_service.last_cleanup_time, tz=UTC))
        except Exception:
            logger.info('error from materialization service', exc_info=True)
            self.send_error(500, reason='error from materialization service')
            return

        try:
            ret = await self.db.materialization.count_documents({'status':{'$in':['waiting','processing']}}, maxTimeMS=1000)
        except Exception:
            logger.info('bad db request', exc_info=True)
            self.send_error(500, reason='bad db request')
            return
        if ret is None:
            self.send_error(500, reason='bad db result')
        else:
            status['num_requests'] = ret
            self.write(status)


class Error(TornadoRequestHandler):
    def prepare(self):
        raise HTTPError(404, 'invalid api route')


@dataclasses.dataclass
class DefaultConfig:
    HOST: str = 'localhost'
    PORT: int = 8080
    DEBUG: bool = False
    OPENID_URL: str = ''
    OPENID_AUDIENCE: str = ''
    ICEPROD_API_ADDRESS: str = 'https://iceprod2-api.icecube.wisc.edu'
    ICEPROD_API_CLIENT_ID: str = ''
    ICEPROD_API_CLIENT_SECRET: str = ''
    DB_URL: str = 'mongodb://localhost/datasets'
    DB_TIMEOUT: int = 60
    DB_WRITE_CONCERN: int = 1
    PROMETHEUS_PORT: int = 0
    CI_TESTING: str = ''


class Server:
    def __init__(self: Self):
        config = from_environment_as_dataclass(DefaultConfig)

        rest_config: dict[str, Any] = {
            'debug': config.DEBUG,
        }
        if config.OPENID_URL:
            logging.info(f'enabling auth via {config.OPENID_URL} for aud "{config.OPENID_AUDIENCE}"')
            rest_config['auth'] = {
                'openid_url': config.OPENID_URL,
                'audience': config.OPENID_AUDIENCE,
            }
        elif config.CI_TESTING:
            rest_config['auth'] = {
                'secret': 'secret',
            }
        else:
            raise RuntimeError('OPENID_URL not specified, and CI_TESTING not enabled!')

        # enable monitoring
        self.prometheus_port = config.PROMETHEUS_PORT if config.PROMETHEUS_PORT > 0 else None
        self.async_monitor = None

        logging_url = config.DB_URL.split('@')[-1] if '@' in config.DB_URL else config.DB_URL
        logging.info(f'DB: {logging_url}')
        db_url, db_name = config.DB_URL.rsplit('/', 1)
        db = motor.motor_asyncio.AsyncIOMotorClient(  # type: ignore[var-annotated]
            db_url,
            timeoutMS=config.DB_TIMEOUT*1000,
            w=config.DB_WRITE_CONCERN,
        )
        logging.info(f'DB name: {db_name}')
        self.db = db[db_name]
        self.indexes = {
            'materialization': {
                'materialization_id_index': {'keys': 'materialization_id', 'unique': True},
                'dataset_id_index': {'keys': 'dataset_id', 'unique': False},
                'status_timestamp_index': {'keys': [('status', pymongo.ASCENDING), ('timestamp', pymongo.ASCENDING)], 'unique': False},
            }
        }

        rest_client: RestClient
        if config.ICEPROD_API_CLIENT_ID and config.ICEPROD_API_CLIENT_SECRET:
            logging.info(f'enabling auth via {config.OPENID_URL} for aud "{config.OPENID_AUDIENCE}"')
            rest_client = ClientCredentialsAuth(
                address=config.ICEPROD_API_ADDRESS,
                token_url=config.OPENID_URL,
                client_id=config.ICEPROD_API_CLIENT_ID,
                client_secret=config.ICEPROD_API_CLIENT_SECRET,
            )
        elif config.CI_TESTING:
            rest_client = RestClient(config.ICEPROD_API_ADDRESS, timeout=1, retries=0)
        else:
            raise RuntimeError('ICEPROD_API_CLIENT_ID or ICEPROD_API_CLIENT_SECRET not specified, and CI_TESTING not enabled!')

        self.materialization_service = MaterializationService(self.db, rest_client)
        self.materialization_service_task = None

        kwargs = IceProdRestConfig(rest_config, database=self.db)
        kwargs['materialization_service'] = self.materialization_service
        kwargs['rest_client'] = rest_client

        server = RestServer(debug=config.DEBUG)

        server.add_route(r'/status/(?P<materialization_id>\w+)', StatusHandler, kwargs)
        server.add_route('/', RequestHandler, kwargs)
        server.add_route(r'/request/(?P<dataset_id>\w+)', RequestDatasetHandler, kwargs)
        server.add_route(r'/request/(?P<dataset_id>\w+)/status', DatasetStatusHandler, kwargs)
        server.add_route('/healthz', HealthHandler, kwargs)
        server.add_route(r'/(.*)', Error)

        server.startup(address=config.HOST, port=config.PORT)

        self.server = server

    async def start(self):
        if self.prometheus_port:
            logging.info("starting prometheus on %r", self.prometheus_port)
            start_http_server(self.prometheus_port)
            i = Info('iceprod', 'IceProd information')
            i.info({
                'version': VERSION_STRING,
                'type': 'materialization',
            })
            self.async_monitor = AsyncMonitor(labels={'type': 'materialization'})
            await self.async_monitor.start()

        for collection in self.indexes:
            existing = await self.db[collection].index_information()
            for name in self.indexes[collection]:
                if name not in existing:
                    logging.info('DB: creating index %s:%s', collection, name)
                    kwargs = self.indexes[collection][name]
                    await self.db[collection].create_index(name=name, **kwargs)

        if not self.materialization_service_task:
            self.materialization_service_task = asyncio.create_task(self.materialization_service.run())

    async def stop(self):
        await self.server.stop()
        if self.materialization_service_task:
            self.materialization_service_task.cancel()
            try:
                await self.materialization_service_task
            except asyncio.CancelledError:
                pass
            finally:
                self.materialization_service_task = None
        if self.async_monitor:
            await self.async_monitor.stop()
