__version__ = '0.5.0'

from json import loads as _loads
from functools import partial as _partial

from aiohttp import ClientSession as _ClientSession, \
    ClientTimeout as _ClientTimeout
from pandas import DataFrame as _DataFrame, to_datetime as _to_dt, \
    to_numeric as _to_num


_YK = ''.maketrans('يك', 'یک')
_DF = _partial(_DataFrame, copy=False)


SESSION : _ClientSession | None = None


class Session(_ClientSession):

    def __init__(self, **kwargs):
        if 'timeout' not in kwargs:
            kwargs['timeout'] = _ClientTimeout(
                total=60, sock_connect=10, sock_read=10)
        super().__init__(**kwargs)

    async def __aenter__(self) -> _ClientSession:
        global SESSION
        SESSION = await super().__aenter__()
        return SESSION


# this function should only be called from _get below
async def _session_get(url: str) -> bytes:
    return await (await SESSION.get(url)).read()


async def _api_json(path) -> list | dict:
    content = await _session_get('https://api.iranetf.org/' + path)
    return _loads(content.decode().translate(_YK))


async def funds() -> _DataFrame:
    j = (await _api_json('odata/company/GetFunds'))['value']
    df = _DF(j)
    df[['UpdateDate', 'CreateDate']] = df[['UpdateDate', 'CreateDate']].apply(_to_dt)
    df['NameDisplay'] = df['NameDisplay'].astype('string', copy=False).str.strip()
    return df


async def fund_portfolio_report_latest(id_: int) -> _DataFrame:
    j = (await _api_json(
        'odata/FundPortfolioReport'
        f'?$top=1'
        f'&$orderby=FromDate desc'
        f'&$filter=CompanyId eq {id_}&$expand=trades'))['value']
    df = _DF(j[0]['Trades'])
    return df


async def funds_deviation_week_month(
    set_index='companyId'
) -> tuple[_DataFrame, _DataFrame]:
    j = await _api_json('bot/funds/fundPriceAndNavDeviation')
    week = _DF(j['seven'])
    month = _DF(j['thirty'])
    if set_index:
        week.set_index(set_index, inplace=True)
        month.set_index(set_index, inplace=True)
    return week, month


async def funds_trade_price(set_index='companyId') -> _DataFrame:
    j = await _api_json('bot/funds/allFundLastStatus/tradePrice')
    df = _DF(j)
    numeric_cols = [
        'tradePrice', 'priceDiff', 'nav', 'navDiff', 'priceAndNavDiff']
    df[numeric_cols] = df[numeric_cols].apply(_to_num, downcast='unsigned')
    if set_index:
        df.set_index(set_index, inplace=True)
    return df


async def fund_trade_info(id_: int | str, month: int) -> _DataFrame:
    j = await _api_json(
        'odata/stockTradeInfo/'
        f'GetCompanyStockTradeInfo(companyId={id_},month={month})')
    df = _DF(j['value'])
    df['Date'] = _to_dt(df['Date'])
    return df


async def companies() -> _DataFrame:
    return _DF((await _api_json('odata/company'))['value'])
