import httpx
from app.workers.base_worker import BaseWorker
from app.job_registry import routing_key
from app.config import settings
from app.workers.utils import resolve_txt

class MTASTSWorker(BaseWorker):
    async def process(self, payload: dict, **meta) -> dict:
        domain = payload.get('domain')
        if not domain:
            raise ValueError('payload.domain is required')
        txt_name = f"_mta-sts.{domain}"
        txts = await resolve_txt(txt_name)
        txt_policy = None
        for t in txts:
            if t.lower().startswith('v=stsv1'):
                txt_policy = t
                break
        # Fetch policy file
        url = f"https://mta-sts.{domain}/.well-known/mta-sts.txt"
        policy = {
            'version': None,
            'mode': None,
            'max_age': None,
            'mx': [],
        }
        content = None
        try:
            timeout = httpx.Timeout(settings.HTTP_TIMEOUT)
            async with httpx.AsyncClient(timeout=timeout) as client:
                resp = await client.get(url)
                if resp.status_code == 200:
                    content = resp.text
        except Exception as e:
            pass
        if content:
            for raw in content.splitlines():
                line = raw.strip()
                if not line or line.startswith('#'):
                    continue
                if ':' in line:
                    k, v = [s.strip() for s in line.split(':', 1)]
                    lk = k.lower()
                    if lk == 'version':
                        policy['version'] = v
                    elif lk == 'mode':
                        policy['mode'] = v
                    elif lk == 'max_age':
                        try:
                            policy['max_age'] = int(v)
                        except ValueError:
                            policy['max_age'] = None
                    elif lk == 'mx':
                        policy['mx'].append(v)
        return {
            'txt_record': txt_policy,
            'policy_url': url,
            'policy': policy,
            'has_policy': bool(policy['version']),
        }

async def main():
    worker = MTASTSWorker(routing_key("MailCheck", "MTASTSCheck"), worker_id="mtasts-worker-1")
    await worker.run()

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