"""Account database operations."""

from dataclasses import asdict
from typing import Literal, overload

import polars as pl

from niveshpy.core.query import ast
from niveshpy.db.database import Database
from niveshpy.db.query import (
    DEFAULT_QUERY_OPTIONS,
    QueryOptions,
    ResultFormat,
    prepare_query_filters,
)
from niveshpy.models.account import AccountRead, AccountWrite


class AccountRepository:
    """Repository for managing investment accounts in the database."""

    _table_name = "accounts"
    _column_mappings = {
        ast.Field.ACCOUNT: ["name", "institution"],
    }

    def __init__(self, db: Database):
        """Initialize the AccountRepository."""
        self._db = db

    def count_accounts(self, options: QueryOptions = DEFAULT_QUERY_OPTIONS) -> int:
        """Count the number of accounts matching the query options."""
        query = f"SELECT COUNT(*) FROM {self._table_name}"
        if options.filters:
            filter_query, params = prepare_query_filters(
                options.filters, self._column_mappings
            )
            query += " WHERE " + filter_query
        else:
            params = ()
        query += ";"

        with self._db.cursor() as cursor:
            res = cursor.execute(query, params)
            count = res.fetchone()
            return count[0] if count is not None else 0

    @overload
    def search_accounts(
        self, options: QueryOptions, format: Literal[ResultFormat.POLARS]
    ) -> pl.DataFrame: ...

    @overload
    def search_accounts(
        self, options: QueryOptions, format: Literal[ResultFormat.SINGLE]
    ) -> tuple | None: ...

    @overload
    def search_accounts(
        self, options: QueryOptions, format: Literal[ResultFormat.LIST]
    ) -> list[tuple]: ...

    def search_accounts(
        self,
        options: QueryOptions = DEFAULT_QUERY_OPTIONS,
        format: ResultFormat = ResultFormat.POLARS,
    ) -> pl.DataFrame | tuple | None | list[tuple]:
        """Search for accounts matching the query options."""
        query = f"SELECT * FROM {self._table_name}"
        if options.filters:
            filter_query, _params = prepare_query_filters(
                options.filters, self._column_mappings
            )
            query += " WHERE " + filter_query
            params = list(_params)
        else:
            params = []

        query += " ORDER BY id"
        if options.limit:
            query += " LIMIT ?"
            params.append(str(options.limit))
        if options.offset:
            query += " OFFSET ?"
            params.append(str(options.offset))
        query += ";"

        with self._db.cursor() as cursor:
            res = cursor.execute(query, params)
            if format == ResultFormat.POLARS:
                return res.pl()
            elif format == ResultFormat.SINGLE:
                return res.fetchone()
            else:
                return res.fetchall()

    def find_account(self, account: AccountWrite) -> AccountRead | None:
        """Find an account by name and institution."""
        query = f"""SELECT * FROM {self._table_name}
                    WHERE name = ? AND institution = ?;"""
        params = (account.name, account.institution)

        with self._db.cursor() as cursor:
            res = cursor.execute(query, params).fetchone()
            return AccountRead(*res) if res else None

    def insert_single_account(self, account: AccountWrite) -> AccountRead | None:
        """Insert a single account into the database."""
        with self._db.cursor() as cursor:
            res = cursor.execute(
                f"""INSERT INTO {self._table_name} (name, institution, metadata)
                VALUES (?, ?, ?)
                RETURNING *;
                """,
                (account.name, account.institution, account.metadata),
            ).fetchone()
            cursor.commit()
            return AccountRead(*res) if res else None

    def get_accounts(self) -> pl.DataFrame:
        """Retrieve all accounts from the database."""
        with self._db.cursor() as cursor:
            return cursor.execute(
                f"SELECT id, name, institution FROM {self._table_name}"
            ).pl()

    def insert_multiple_accounts(
        self, accounts: list[AccountWrite]
    ) -> list[AccountRead]:
        """Insert multiple accounts into the database."""
        with self._db.cursor() as cursor:
            cursor.begin()
            cursor.register("new_accounts", pl.from_dicts(map(asdict, accounts)))
            cursor.execute(
                f"""MERGE INTO {self._table_name} target
                USING (SELECT * FROM new_accounts) AS new
                ON target.name = new.name AND target.institution = new.institution
                WHEN MATCHED THEN UPDATE SET metadata = new.metadata
                WHEN NOT MATCHED THEN INSERT BY NAME
                RETURNING *;
                """
            )
            res = cursor.fetchall()
            cursor.commit()
            return [AccountRead(*data) for data in res]

    # def add_accounts(self, accounts: Iterable[AccountWrite]) -> Iterable[AccountRead]:
    #     """Add new accounts to the database."""
    #     with self._db.cursor() as cursor:
    #         cursor.register("new_accounts", pl.from_dicts(map(asdict, accounts)))
    #         data = cursor.execute(
    #             f"""MERGE INTO {self._table_name} target
    #             USING (SELECT * FROM new_accounts) AS new
    #             ON target.name = new.name AND target.institution = new.institution
    #             WHEN NOT MATCHED THEN INSERT BY NAME;

    #             FROM {self._table_name}
    #             ORDER BY id;
    #             """
    #         )
    #         cursor.commit()
    #         return starmap(AccountRead, data.fetchall())

    def get_account(self, id: int) -> AccountRead | None:
        """Retrieve an account by its ID."""
        query = f"SELECT id, name, institution, created_at, metadata FROM {self._table_name} WHERE id = ?;"
        params = (id,)

        with self._db.cursor() as cursor:
            res = cursor.execute(query, params).fetchone()
            return AccountRead(*res) if res else None

    def delete_account(self, id: int) -> AccountRead | None:
        """Delete an account by its ID.

        Returns the deleted account if successful, None otherwise.
        """
        with self._db.cursor() as cursor:
            res = cursor.execute(
                f"DELETE FROM {self._table_name} WHERE id = ? RETURNING *;",
                (id,),
            )
            cursor.commit()
            data = res.fetchone()
            return AccountRead(*data) if data else None
