"""
Account management for the Tradier API.

This module provides the Account class for retrieving account information,
balances, and positions.
"""

from __future__ import annotations

from datetime import datetime
from decimal import Decimal
from enum import Enum
from typing import TYPE_CHECKING

from loguru import logger
from pydantic import Field

from optrabot.broker.tradier.utils import TradierData

if TYPE_CHECKING:
    from optrabot.broker.tradier.session import Session


# ========== Balance Models ==========


class MarginBalance(TradierData):
    """Margin-specific balance information."""

    fed_call: Decimal = Decimal('0')
    maintenance_call: Decimal = Decimal('0')
    option_buying_power: Decimal = Decimal('0')
    stock_buying_power: Decimal = Decimal('0')
    stock_short_value: Decimal = Decimal('0')
    sweep: Decimal = Decimal('0')


class CashBalance(TradierData):
    """Cash-specific balance information."""

    cash_available: Decimal = Decimal('0')
    sweep: Decimal = Decimal('0')
    unsettled_funds: Decimal = Decimal('0')


class PDTBalance(TradierData):
    """Pattern Day Trader balance information."""

    fed_call: Decimal = Decimal('0')
    maintenance_call: Decimal = Decimal('0')
    option_buying_power: Decimal = Decimal('0')
    stock_buying_power: Decimal = Decimal('0')
    stock_short_value: Decimal = Decimal('0')


class AccountBalance(TradierData):
    """
    Account balance information.

    Contains the current balance, equity, and buying power for an account.
    Different account types (margin, cash, pdt) have different sub-fields.

    Example::

        balance = account.get_balance(session)
        print(f"Total Equity: ${balance.total_equity}")
        print(f"Cash: ${balance.total_cash}")
        if balance.margin:
            print(f"Buying Power: ${balance.margin.stock_buying_power}")
    """

    # Core balance fields
    account_number: str
    account_type: str
    total_equity: Decimal = Decimal('0')
    total_cash: Decimal = Decimal('0')
    equity: Decimal = Decimal('0')
    market_value: Decimal = Decimal('0')

    # Position values
    long_market_value: Decimal = Decimal('0')
    short_market_value: Decimal = Decimal('0')
    stock_long_value: Decimal = Decimal('0')
    option_long_value: Decimal = Decimal('0')
    option_short_value: Decimal = Decimal('0')

    # P&L
    open_pl: Decimal = Decimal('0')
    close_pl: Decimal = Decimal('0')

    # Margin requirements
    current_requirement: Decimal = Decimal('0')
    option_requirement: Decimal = Decimal('0')

    # Cash status
    pending_cash: Decimal = Decimal('0')
    uncleared_funds: Decimal = Decimal('0')
    pending_orders_count: int = 0

    # Type-specific balances (only one will be populated based on account type)
    margin: MarginBalance | None = None
    cash: CashBalance | None = None
    pdt: PDTBalance | None = None

    def __repr__(self) -> str:
        return (
            f'AccountBalance(account={self.account_number!r}, '
            f'equity=${self.total_equity}, cash=${self.total_cash})'
        )


class AccountClassification(str, Enum):
    """Account classification types."""

    INDIVIDUAL = 'individual'
    INDIVIDUAL_CASH = 'individual_cash'
    INDIVIDUAL_MARGIN = 'individual_margin'
    ENTITY = 'entity'
    ENTITY_CASH = 'entity_cash'
    ENTITY_MARGIN = 'entity_margin'
    JOINT = 'joint'
    JOINT_CASH_SURVIVOR = 'joint_cash_survivor'
    JOINT_MARGIN_SURVIVOR = 'joint_margin_survivor'
    JOINT_MARGIN_TENANT = 'joint_margin_tenant'
    CORPORATE = 'corporate'
    TRADITIONAL_IRA = 'traditional_ira'
    ROTH_IRA = 'roth_ira'
    ROLLOVER_IRA = 'rollover_ira'
    SEP_IRA = 'sep_ira'
    TRADITIONAL_SEP_IRA = 'traditional_sep_ira'
    CUSTODIAL_CASH = 'custodial_cash'


class AccountStatus(str, Enum):
    """Account status types."""

    ACTIVE = 'active'
    CLOSED = 'closed'


class AccountType(str, Enum):
    """Account types."""

    CASH = 'cash'
    MARGIN = 'margin'


class Account(TradierData):
    """
    Represents a Tradier brokerage account.

    Use the class methods to retrieve accounts from the API:

    Example::

        from optrabot.broker.tradier import Session, Account

        session = Session("your-access-token")

        # Get all accounts
        accounts = Account.get_accounts(session)

        # Get a specific account
        account = Account.get(session, "VA000001")
    """

    account_number: str
    classification: str
    date_created: datetime
    day_trader: bool
    option_level: int
    status: AccountStatus
    type: AccountType = Field(alias='type')
    last_update_date: datetime

    def __repr__(self) -> str:
        return (
            f'Account(account_number={self.account_number!r}, '
            f'type={self.type!r}, status={self.status!r})'
        )

    # ========== Class Methods for Retrieving Accounts ==========

    @classmethod
    def get_accounts(cls, session: Session) -> list[Account]:
        """
        Get all accounts associated with the current user.

        :param session: The session to use for the request.
        :returns: A list of Account objects.

        Example::

            accounts = Account.get_accounts(session)
            for account in accounts:
                print(f"{account.account_number}: {account.type}")
        """
        data = session._get('/user/profile')
        return cls._parse_accounts(data)

    @classmethod
    async def a_get_accounts(cls, session: Session) -> list[Account]:
        """
        Get all accounts associated with the current user (async).

        :param session: The session to use for the request.
        :returns: A list of Account objects.

        Example::

            accounts = await Account.a_get_accounts(session)
            for account in accounts:
                print(f"{account.account_number}: {account.type}")
        """
        data = await session._a_get('/user/profile')
        return cls._parse_accounts(data)

    @classmethod
    def get(cls, session: Session, account_number: str) -> Account:
        """
        Get a specific account by account number.

        :param session: The session to use for the request.
        :param account_number: The account number to retrieve.
        :returns: The Account object.
        :raises TradierError: If the account is not found.

        Example::

            account = Account.get(session, "VA000001")
            print(f"Option Level: {account.option_level}")
        """
        accounts = cls.get_accounts(session)
        return cls._find_account(accounts, account_number)

    @classmethod
    async def a_get(cls, session: Session, account_number: str) -> Account:
        """
        Get a specific account by account number (async).

        :param session: The session to use for the request.
        :param account_number: The account number to retrieve.
        :returns: The Account object.
        :raises TradierError: If the account is not found.

        Example::

            account = await Account.a_get(session, "VA000001")
            print(f"Option Level: {account.option_level}")
        """
        accounts = await cls.a_get_accounts(session)
        return cls._find_account(accounts, account_number)

    # ========== Instance Methods for Account Data ==========

    def get_balance(self, session: Session) -> AccountBalance:
        """
        Get the current balance for this account.

        :param session: The session to use for the request.
        :returns: An AccountBalance object with balance details.

        Example::

            balance = account.get_balance(session)
            print(f"Total Equity: ${balance.total_equity}")
            print(f"Total Cash: ${balance.total_cash}")
            if balance.margin:
                print(f"Stock Buying Power: ${balance.margin.stock_buying_power}")
        """
        data = session._get(f'/accounts/{self.account_number}/balances')
        return self._parse_balance(data)

    async def a_get_balance(self, session: Session) -> AccountBalance:
        """
        Get the current balance for this account (async).

        :param session: The session to use for the request.
        :returns: An AccountBalance object with balance details.

        Example::

            balance = await account.a_get_balance(session)
            print(f"Total Equity: ${balance.total_equity}")
        """
        data = await session._a_get(f'/accounts/{self.account_number}/balances')
        return self._parse_balance(data)

    def _parse_balance(self, data: dict) -> AccountBalance:
        """Parse the balance response and return an AccountBalance object."""
        balances_data = data.get('balances', {})
        logger.debug(
            'Retrieved balance for account {}: equity=${}',
            self.account_number,
            balances_data.get('total_equity', 0),
        )
        return AccountBalance(**balances_data)

    # ========== Private Helper Methods ==========

    @classmethod
    def _parse_accounts(cls, data: dict) -> list[Account]:
        """Parse the profile response and return a list of accounts."""
        profile = data.get('profile', {})
        accounts_data = profile.get('account', [])

        # Handle single account (API returns dict) vs multiple (returns list)
        if isinstance(accounts_data, dict):
            accounts_data = [accounts_data]

        accounts = [cls(**account) for account in accounts_data]
        logger.trace('Retrieved {} account(s)', len(accounts))
        return accounts

    @classmethod
    def _find_account(cls, accounts: list[Account], account_number: str) -> Account:
        """Find an account by account number."""
        from optrabot.broker.tradier.exceptions import TradierError

        for account in accounts:
            if account.account_number == account_number:
                return account

        raise TradierError(f'Account {account_number} not found')
