import yaml
import functools
import itertools
import numpy as np
import pandas as pd
from sqlalchemy import create_engine


class Render:
    def __init__(self, config_file: str, trino_host: str):
        self._ctx_stack = []
        with open(config_file, "r") as stream:
            self.configs = yaml.safe_load(stream)
        self.engine = create_engine(f'trino://admin@{trino_host}:8080')

    def __enter__(self):
        return self

    def __exit__(self, *args, **kwargs):
        self._ctx_stack.pop()

    def context(self, **ctx):
        self._ctx_stack.append(ctx)
        return self

    @functools.lru_cache
    def query(self, sql):
        return pd.read_sql(sql, self.engine)

    def dataframe(self, cfg_id):
        assert cfg_id in self.configs, f"{cfg_id} not found"
        sql = self.configs[cfg_id].format(**self.ctx)
        return self.query(sql)

    def __getattr__(self, cfg_id):
        if cfg_id == 'ctx':
            return dict(itertools.chain(*[c.items() for c in self._ctx_stack]))

        df = self.dataframe(cfg_id)
        if len(df) == 0:
            return np.nan
        if len(df) != 1 or len(df.columns) != 1:
            return df.to_markdown(index=False)
        if df.iloc[0, 0] is None:
            return np.nan
        return df.iloc[0, 0]