import builtins
from collections.abc import Sequence
from pathlib import Path
from typing import Optional, Union
from uuid import UUID

import srsly
from radicli import Arg
from wasabi import msg

from prodigy_teams_pam_sdk.models import DatasetDetail, DatasetSummary

from ..cli import cli
from ..errors import CLIError, ProdigyTeamsErrors
from ..messages import Messages
from ..query import collect_from_pages, resolve_dataset, resolve_dataset_id
from ..ui import (
    PageProgress,
    print_info_table,
    print_mutation_result,
    print_table_with_select,
)
from ._state import get_auth_state


@cli.subcommand(
    "datasets",
    "list",
    # fmt: off
    select=Arg(
        "--select",
        help=Messages.select.format(opts=builtins.list(DatasetSummary.model_fields)),
    ),
    as_json=Arg("--json", help=Messages.as_json),
    # fmt: on
)
def list(
    select: builtins.list[str] = ["id", "name", "kind"], as_json: bool = False
) -> Sequence[DatasetSummary]:
    """List all Datasets"""
    client = get_auth_state().client
    items = builtins.list(client.dataset.all())
    print_table_with_select(items, select=select, as_json=as_json)
    return items


@cli.subcommand(
    "datasets",
    "info",
    name_or_id=Arg(help=Messages.name_or_id.format(noun="dataset")),
    cluster_id=Arg(help=Messages.cluster_id.format(noun="dataset")),
    select=Arg(
        "--select",
        help=Messages.select.format(opts=builtins.list(DatasetDetail.model_fields)),
    ),
    as_json=Arg("--json", help=Messages.as_json),
)
def info(
    name_or_id: Union[str, UUID],
    cluster_id: Optional[UUID] = None,
    select: Optional[builtins.list[str]] = None,
    as_json: bool = False,
) -> DatasetDetail:
    """Get detailed info for a Dataset"""
    res = resolve_dataset(name_or_id, broker_id=cluster_id)
    print_info_table(res, as_json=as_json, select=select)
    return res


@cli.subcommand(
    "datasets",
    "create",
    name=Arg(help=Messages.name.format(noun="dataset")),
    kind=Arg("--kind", help=Messages.kind.format(noun="dataset")),
    exists_ok=Arg("--exists-ok", help=Messages.exists_ok),
    as_json=Arg("--json", help=Messages.as_json),
)
def create(
    name: str, kind: str, exists_ok: bool = False, as_json: bool = False
) -> DatasetSummary | DatasetDetail | None:
    """Create a new dataset"""
    auth = get_auth_state()
    client = auth.client
    broker_id = auth.broker_id
    try:
        res = client.dataset.create(name=name, kind=kind, broker_id=broker_id)
    except ProdigyTeamsErrors.DatasetExists:
        if exists_ok:
            msg.info(Messages.T001.format(noun="dataset", name=name))
            return None
        raise CLIError(Messages.E002.format(noun="dataset", name=name))
    except ProdigyTeamsErrors.DatasetInvalid:
        raise CLIError(Messages.E004.format(noun="dataset", name=name))
    except ProdigyTeamsErrors.DatasetForbiddenCreate:
        raise CLIError(Messages.E003.format(noun="dataset", name=name))
    print_mutation_result(
        {"id": str(res.id), "name": res.name},
        Messages.T002.format(noun="dataset", name=res.name),
        as_json=as_json,
    )
    if not as_json:
        print_info_table(res)
    return res


@cli.subcommand(
    "datasets",
    "delete",
    name_or_id=Arg(help=Messages.name_or_id.format(noun="dataset")),
    cluster_id=Arg(help=Messages.cluster_id.format(noun="dataset")),
    as_json=Arg("--json", help=Messages.as_json),
)
def delete(
    name_or_id: Union[str, UUID],
    cluster_id: Optional[UUID] = None,
    as_json: bool = False,
) -> UUID:
    """Delete a dataset"""
    dataset_id = resolve_dataset_id(name_or_id, broker_id=cluster_id)
    auth = get_auth_state()
    try:
        auth.client.dataset.delete(id=dataset_id)
    except (
        ProdigyTeamsErrors.ProjectForbiddenDelete,
        ProdigyTeamsErrors.ProjectNotFound,
    ):
        raise CLIError(Messages.E006.format(noun="dataset", name=name_or_id))
    else:
        print_mutation_result(
            {"id": str(dataset_id), "deleted": True},
            Messages.T003.format(noun="dataset", name=name_or_id),
            as_json=as_json,
        )
    return dataset_id


@cli.subcommand(
    "datasets",
    "export",
    name_or_id=Arg(help=Messages.name_or_id.format(noun="dataset")),
    output_file=Arg("--output", "-o", help=Messages.export_output),
)
def export(
    name_or_id: Union[str, UUID],
    output_file: Union[str, Path] = "-",
) -> None:
    """Export all the examples from a dataset and save it in the designated file as
    JSONL (newline-delimited JSON).
    """
    auth = get_auth_state()
    cluster_id = auth.broker_id

    dataset = resolve_dataset(name_or_id, broker_id=cluster_id)

    def get_examples_page(page: int):
        return auth.broker_client.datasets.read_examples(
            [dataset.name],
            page=page,
        )

    with PageProgress(collect_from_pages(get_examples_page)) as pages:
        srsly.write_jsonl(
            output_file,
            (example.content for page in pages for example in page.items),
            append_new_line=False,
        )
