from typing import Iterator

from prometheus_client.metrics_core import GaugeMetricFamily, CounterMetricFamily, Metric
from sqlalchemy import text

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


class DefinedProxiesCollector:
    def collect(self) -> Iterator[Metric]:
        with app.app_context():
            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'])  # type: ignore[no-untyped-call]
            for row in result:
                c.add_metric([str(row[0]), row[1], row[2], str(row[3]), row[4]],
                             row[5])  # type: ignore[no-untyped-call]
            yield c


class BlockedProxiesCollector:
    def collect(self) -> Iterator[Metric]:
        with app.app_context():
            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'])  # type: ignore[no-untyped-call]
                for row in result:
                    c.add_metric([str(row[0]), row[1], row[2], str(row[3]), row[4], row[5]],
                                 row[6])  # type: ignore[no-untyped-call]
                yield c


class AutomationCollector:
    def collect(self) -> Iterator[Metric]:
        with app.app_context():
            c = GaugeMetricFamily("automation_state", "The automation state (0: idle, 1: running, 2: error)",
                                  labels=['automation_name'])  # type: ignore[no-untyped-call]
            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)  # type: ignore[no-untyped-call]
                elif automation.state == AutomationState.RUNNING:
                    c.add_metric([automation.short_name], 1)  # type: ignore[no-untyped-call]
                else:
                    c.add_metric([automation.short_name], 2)  # type: ignore[no-untyped-call]
            yield c