# Copyright (c) 2024 Julien Peloton
#
# This file is part of TOM Toolkit
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.
import time
import logging
import traceback

from fink_client.consumer import AlertConsumer

from tom_alertstreams.alertstreams.alertstream import AlertStream
from tom_targets.models import Target, TargetList

from psycopg2.errors import UniqueViolation
from django.db.utils import IntegrityError as DJ_IntegrityError
from sqlite3 import IntegrityError as SQL_IntegrityError
from django.contrib.auth.models import Group
from guardian.shortcuts import assign_perm

logger = logging.getLogger(__name__)


class FinkAlertStream(AlertStream):
    """Poll alerts from a stream generated by Fink"""

    required_keys = [
        "URL",
        "GROUP_ID",
        "USERNAME",
        "TOPIC",
        "TOPIC_HANDLERS",
        "MAX_POLL_NUMBER",
        "TIMEOUT",
    ]
    allowed_keys = [
        "URL",
        "GROUP_ID",
        "USERNAME",
        "PASSWORD",
        "TOPIC",
        "TOPIC_HANDLERS",
        "MAX_POLL_NUMBER",
        "TIMEOUT",
    ]

    def __init__(self, *args, **kwargs) -> None:
        """Initialise credentials and target list"""
        super().__init__(*args, **kwargs)

        # Define target list based on topics
        public_group, _ = Group.objects.get_or_create(name="Public")
        tl, is_created = TargetList.objects.get_or_create(name=self.topic)
        assign_perm("tom_targets.view_targetlist", public_group, tl)

    def listen(self):
        """Listen to one topic from Fink Kafka server

        Each alert is saved into the target list named from
        the topic.

        Note
        ----------
        This only handles one topic at once for the moment.
        """
        super().listen()

        myconfig = {
            "username": self.username,
            "bootstrap.servers": self.url,
            "group.id": self.group_id,
        }

        consumer = AlertConsumer([self.topic], myconfig, schema_path=None)

        poll_number = 0
        header = "FinkAlertStream.listen"
        while poll_number < int(self.max_poll_number):
            try:
                logger.info(
                    f"{header} opening stream: {self.url} with group.id: {self.group_id} (call number: {poll_number})"
                )
                topic, alert, key = consumer.poll(timeout=int(self.timeout))

                if topic is not None:
                    # TODO: handle MMA vs regular streams
                    self.alert_handler["fink.stream"](alert, topic)
                else:
                    logger.info("No alerts received")
                poll_number += 1
            except Exception as ex:
                logger.error(f"{header}: {ex}")
                logger.error(traceback.format_exc())
                break
        consumer.close()


def alert_logger(alert, topic):
    """Basic alert handler for Fink

    This alert handler simply display on screen basic information,
    and save the alert as a new Target.

    Parameters
    ----------
    alert: dic
        Dictionary containing alert data. See `consumer.poll`.
    topic: str
        Topic name

    Warnings
    ----------
    UniqueViolation, SQL_IntegrityError, DJ_IntegrityError
        If the target is already saved

    Raises
    ------
    Exception (base)
        for any other failures than name clash when
        saving the target in the database.

    """
    utc = time.strftime("%Y-%m-%d %H:%M:%S", time.gmtime())
    logger.info(f"fink.alert_logger topic: {topic}")
    logger.info(
        "fink.alert_logger value: {} emitted {} JD (received {})".format(
            alert["objectId"], alert["candidate"]["jd"], utc
        )
    )

    mytarget = Target(
        name=alert["objectId"],
        type="SIDEREAL",
        ra=alert["candidate"]["ra"],
        dec=alert["candidate"]["dec"],
        epoch=alert["candidate"]["jd"],
    )

    try:
        mytarget.save(
            extras={
                "fink broker link": "https://fink-portal.org/{}".format(
                    alert["objectId"]
                )
            }
        )
    except (UniqueViolation, SQL_IntegrityError, DJ_IntegrityError):
        logger.warning(f"Target {mytarget} already in the database")
        pass
    except Exception:
        logger.error("error when trying to save new alerts in the db", exc_info=1)
        logger.error(traceback.format_exc())
