import asyncio
from typing import Dict, Any, List
from sqlalchemy import select
from app.db import SessionLocal
from app.models import Task
from app.workers.base_worker import BaseWorker
from app.job_registry import routing_key
from app.workers.utils import resolve_mx, smtp_starttls_probe

class TLSCheckWorker(BaseWorker):
    async def process(self, payload: dict, **meta) -> dict:
        domain = payload.get("domain")
        job_id = meta.get('job_id')
        if not domain:
            raise ValueError("payload.domain is required")
        # Versuche MX aus MXValidation-Resultat zu lesen
        mx_hosts = []
        if job_id:
            async with SessionLocal() as session:
                res = await session.execute(select(Task).where(Task.job_id == job_id, Task.name == 'MXValidation'))
                mx_task = res.scalar_one_or_none()
                if mx_task and mx_task.result and 'mx_records' in mx_task.result:
                    mx_hosts = [r['exchange'] for r in mx_task.result.get('mx_records', [])]
        if not mx_hosts:
            mx = await resolve_mx(domain)
            mx_hosts = [m.exchange for m in mx]
        if not mx_hosts:
            return {"mx_hosts": [], "summary": {"total": 0, "starttls_supported": 0, "tls_ok": 0}}
        results: List[Dict[str, Any]] = []
        for host in mx_hosts:
            probe = await smtp_starttls_probe(host, 25)
            results.append(probe)
        starttls_supported = sum(1 for r in results if r.get('starttls_advertised'))
        tls_ok = sum(1 for r in results if r.get('tls_established'))
        return {
            "mx_hosts": results,
            "summary": {
                "total": len(results),
                "starttls_supported": starttls_supported,
                "tls_ok": tls_ok,
            },
        }

async def main():
    worker = TLSCheckWorker(routing_key("MailCheck", "TLSCheck"), worker_id="tls-worker-1")
    await worker.run()

if __name__ == "__main__":
    import asyncio as _a
    _a.run(main())