"""
DGA-Logit: TF-IDF + Lexical Features + Logistic Regression for DGA detection.
Trained on 54 DGA families (~845K samples).
artifacts.joblib contains: {'model': LR, 'vectorizer': TF-IDF, 'scaler': StandardScaler}
"""
import re
import math
import numpy as np
import pandas as pd
from collections import Counter
from urllib.parse import urlparse

VOWELS = set('aeiou')
CONSONANTS = set('bcdfghjklmnpqrstvwxyz')
DOMAIN_RE = re.compile(r'^[a-z0-9.-]+$')
LETTER_RE = re.compile(r'[a-z]')
DIGIT_RE = re.compile(r'[0-9]')


def _hostname_from_url(value: str) -> str:
    parsed = urlparse(value)
    if parsed.netloc:
        return parsed.netloc
    if parsed.scheme and parsed.path:
        return parsed.path
    return value


def normalize_domain(value: str) -> str:
    domain = str(value).strip().lower()
    domain = _hostname_from_url(domain)
    domain = domain.split('@')[-1].split('/')[0].split(':')[0].rstrip('.')
    domain = re.sub(r'\s+', '', domain)
    if not DOMAIN_RE.match(domain):
        domain = re.sub(r'[^a-z0-9.-]', '', domain)
    return domain


def shannon_entropy(value: str) -> float:
    if not value:
        return 0.0
    counts = Counter(value)
    total = len(value)
    return -sum((n / total) * math.log2(n / total) for n in counts.values())


def _max_run(value: str, matcher) -> int:
    best = current = 0
    for ch in value:
        if matcher.match(ch):
            current += 1
            best = max(best, current)
        else:
            current = 0
    return best


def _split_parts(domain: str):
    parts = [p for p in domain.split('.') if p]
    if not parts:
        return '', ''
    sld = parts[-2] if len(parts) >= 2 else parts[-1]
    return sld, parts[-1]


def _extract_lexical_features(domains: pd.Series) -> np.ndarray:
    rows = []
    for value in domains:
        domain = normalize_domain(value)
        sld, tld = _split_parts(domain)
        letters = [c for c in domain if c.isalpha()]
        digits = [c for c in domain if c.isdigit()]
        chars = [c for c in domain if c.isalnum()]
        vowel_count = sum(1 for c in letters if c in VOWELS)
        consonant_count = sum(1 for c in letters if c in CONSONANTS)
        length = max(len(domain), 1)
        rows.append([
            len(domain),
            len(sld),
            len(tld),
            max(domain.count('.') - 1, 0),
            len(digits) / length,
            vowel_count / length,
            consonant_count / length,
            (len(set(chars)) / max(len(chars), 1)) if chars else 0.0,
            domain.count('-'),
            domain.count('.'),
            _max_run(domain, DIGIT_RE),
            _max_run(domain, LETTER_RE),
            shannon_entropy(domain),
            float(domain[:1].isdigit()),
            float(domain[-1:].isdigit()),
        ])
    return np.asarray(rows, dtype=float)


def load_model(artifacts_path: str):
    """Load artifacts dict from joblib file."""
    import joblib
    return joblib.load(artifacts_path)


def predict(artifacts, domains):
    """
    Predict DGA vs legit for a list of domain strings.
    artifacts: dict with keys 'model', 'vectorizer', 'scaler'
    Returns list of dicts: [{"domain": ..., "label": "dga"/"legit", "score": float}]
    """
    from scipy import sparse

    if isinstance(domains, str):
        domains = [domains]

    model = artifacts['model']
    vectorizer = artifacts['vectorizer']
    scaler = artifacts['scaler']

    series = pd.Series(domains)
    domains_norm = series.map(normalize_domain)
    X_tfidf = vectorizer.transform(domains_norm)
    X_lex = scaler.transform(_extract_lexical_features(series))
    X = sparse.hstack([X_tfidf, sparse.csr_matrix(X_lex)], format='csr')

    scores = model.predict_proba(X)[:, 1]
    preds = (scores >= 0.5).astype(int)

    return [
        {"domain": d, "label": "dga" if p == 1 else "legit", "score": round(float(s), 4)}
        for d, p, s in zip(domains, preds, scores)
    ]