from abc import abstractmethod
from typing import Iterator, Optional

from flask import Flask
from prometheus_client.metrics_core import GaugeMetricFamily, CounterMetricFamily, Metric
from prometheus_client.registry import Collector
from sqlalchemy import text

from app.extensions import db
from app.models.automation import Automation, AutomationState


class FlaskCollector(Collector):
    _app: Optional[Flask]

    def __init__(self, app: Optional[Flask] = None) -> None:
        self._app = app

    def init_app(self, app: Flask) -> None:
        self._app = app

    def collect(self) -> Iterator[Metric]:
        if self._app:
            with self._app.app_context():
                return self.collect_in_ctx(self._app)
        raise RuntimeError("Flask collector was not initialised with app")

    @abstractmethod
    def collect_in_ctx(self, app: Flask) -> Iterator[Metric]:
        raise NotImplementedError()


class DefinedProxiesCollector(FlaskCollector):
    def collect_in_ctx(self, app: Flask) -> Iterator[Metric]:
        conn = db.engine.connect()
        result = conn.execute(text("""
                SELECT origin.group_id, "group".group_name, proxy.provider, proxy.pool_id, pool.pool_name,
                       COUNT(proxy.id) FROM proxy, origin, pool, "group"
                WHERE proxy.origin_id = origin.id
                    AND origin.group_id = "group".id
                    AND proxy.pool_id = pool.id
                    AND proxy.destroyed IS NULL
                GROUP BY origin.group_id, "group".group_name, proxy.provider, proxy.pool_id, pool.pool_name;
                """))
        c = GaugeMetricFamily("defined_proxies", "Number of proxies currently defined for deployment",
                              labels=['group_id', 'group_name', 'provider', 'pool_id',
                                      'pool_name'])
        for row in result:
            c.add_metric([str(row[0]), row[1], row[2], str(row[3]), row[4]],
                         row[5])
        yield c


class BlockedProxiesCollector(FlaskCollector):
    def collect_in_ctx(self, app: Flask) -> Iterator[Metric]:
        with db.engine.connect() as conn:
            result = conn.execute(text("""
                    SELECT origin.group_id, "group".group_name, proxy.provider, proxy.pool_id, pool.pool_name,
                           proxy.deprecation_reason, COUNT(proxy.id) FROM proxy, origin, pool, "group"
                    WHERE proxy.origin_id = origin.id
                        AND origin.group_id = "group".id
                        AND proxy.pool_id = pool.id
                        AND proxy.deprecated IS NOT NULL
            GROUP BY origin.group_id, "group".group_name, proxy.provider, proxy.pool_id, pool.pool_name,
                     proxy.deprecation_reason;
                    """))
            c = CounterMetricFamily("deprecated_proxies",
                                    "Number of proxies deprecated",
                                    labels=['group_id', 'group_name', 'provider', 'pool_id', 'pool_name',
                                            'deprecation_reason'])
            for row in result:
                c.add_metric([str(row[0]), row[1], row[2], str(row[3]), row[4], row[5]],
                             row[6])
            yield c


class AutomationCollector(FlaskCollector):
    def collect_in_ctx(self, app: Flask) -> Iterator[Metric]:
        c = GaugeMetricFamily("automation_state", "The automation state (0: idle, 1: running, 2: error)",
                              labels=['automation_name'])
        automations = Automation.query.all()
        for automation in automations:
            if automation.short_name in app.config['HIDDEN_AUTOMATIONS']:
                continue
            if automation.state == AutomationState.IDLE:
                c.add_metric([automation.short_name], 0)
            elif automation.state == AutomationState.RUNNING:
                c.add_metric([automation.short_name], 1)
            else:
                c.add_metric([automation.short_name], 2)
        yield c