from ts.client import TradeStationClient
from api.model import Positions
# from private.config import *


class TradeStationConnector:
    field_conversion_map = dict(
        financialCurrency='currency',
        type='quoteType',
        current_annual_dividend_yield='dividendYield',
        current_annual_dividend_rate='dividendRate',
        trailing_pe='trailingPE',
        forward_pe='forwardPE'
    )

    def __init__(self, secrets: dict):
        self.ts_client = TradeStationClient(
            username=secrets['username'],
            client_id=secrets['client_id'],
            client_secret=secrets['client_secret'],
            redirect_uri=secrets['redirect_uri'],
            paper_trading=False
        )
        self.ts_client.login()

    @staticmethod
    def to_camel_case(snake_str):
        # TODO move to base class
        components = snake_str.split('_')
        # We capitalize the first letter of each component with the 'title' method and join them together.
        return ''.join(x.title() for x in components)

    def convert_model_fields(self, position_data: dict):
        # TODO move to base class
        security_record = {}
        for field in Positions.__table__.columns.keys():
            converted_camel_case = self.to_camel_case(field)
            if converted_camel_case in position_data:
                converted_field = converted_camel_case
            elif field in position_data:
                converted_field = field
            elif field in self.field_conversion_map:
                converted_field = self.field_conversion_map[field]
            else:
                continue
            security_record[field] = position_data[converted_field]
        return security_record

    def get_user_accounts(self):
        accounts = self.ts_client.user_accounts(self.ts_client.config['username'])
        return accounts

    def get_account_balance(self, account_key=None):
        accounts = self.get_user_accounts()
        balance = self.ts_client.account_balances([str(accounts[0]['Key'])])
        return balance

    def get_account_positions(self, account_key=None) -> dict:
        accounts = self.get_user_accounts()
        positions = self.ts_client.account_positions([str(accounts[0]['Key'])], None)
        converted_positions = []
        for position in positions:
            converted_positions.append(self.convert_model_fields(position))
        return converted_positions
