import pandas as pd
from loguru import logger

from ..db.session import get_db
from ..db.queries import QueryLoader
from ..config.paths import QUERY_PATHS, PATHS
from ..config.countries import COUNTRY_DICT, LABEL_DICT
from ..config.email import MULTI_MARKET_REPORTS


class DistroService:
    def __init__(self):
        self.db = get_db()
        self.query_loader = QueryLoader()
    
    def update_distro(self) -> None:
        logger.info("Updating distribution list from DB")
        
        query = self.query_loader.load(QUERY_PATHS.distro)
        distro_df = self.db.query(query)
        distro_df.to_csv(PATHS.distro_csv, index=False)
        
        logger.success(f"Distro saved to {PATHS.distro_csv}")
    
    def get_distro(self, mode: str = "test") -> dict[str, list[str]]:
        distro_df = pd.read_csv(PATHS.distro_csv)
        distro_df = distro_df[distro_df["mode"] == mode]
        
        distro_dict: dict[str, list[str]] = {}
        
        for cc in COUNTRY_DICT:
            distro_dict[cc] = distro_df[distro_df["country_code"] == cc]["email"].tolist()
        
        for label in LABEL_DICT:
            distro_dict[label] = distro_df[distro_df["country_code"] == label]["email"].tolist()
        
        for cc, info in COUNTRY_DICT.items():
            if not info.get("generate_ro_mail"):
                continue
            distro_dict[f"x{cc}"] = distro_df[distro_df["country_code"] == f"x{cc}"]["email"].tolist()
        
        distro_dict["CEA"] = distro_df[distro_df["country_code"] == "CEA"]["email"].tolist()
        
        for report in MULTI_MARKET_REPORTS:
            key = report["distro_key"]
            distro_dict[key] = distro_df[distro_df["country_code"] == key]["email"].tolist()
        
        all_recipients = distro_df[distro_df["country_code"] == "ALL"]["email"].tolist()
        for key in distro_dict:
            distro_dict[key] += all_recipients
        
        return distro_dict