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

from radicli import Arg
from wasabi import msg

from ..cli import cli
from ..errors import BrokerError, CLIError, HTTPXErrors, ProdigyTeamsError
from ..messages import Messages
from ..prodigy_teams_pam_sdk.models import AgentDetail, AgentSummary, RecipeDetail
from ..query import (
    delete_job,
    resolve_agent,
    resolve_recipe,
    resolve_task,
    start_job,
    stop_job,
)
from ..ui import (
    print_args_table,
    print_info_table,
    print_logs,
    print_recipes_help,
    print_table_with_select,
)
from ._recipe_subcommand import (
    create_from_recipe,
    create_from_recipe_file,
    request_recipes,
)
from ._state import get_auth_state, get_saved_settings

# Module-level holder for recipe schema passed from create() to start().
_pending_recipe_schema: RecipeDetail | None = None


@cli.subcommand_with_extra(
    "agents",
    "create",
    exists_ok=Arg("--exists-ok", help=Messages.exists_ok),
    no_start=Arg("--no-start", help=Messages.no_start),
    no_wait=Arg("--no-wait", help=Messages.no_wait.format(noun="agent")),
    _show_help=Arg("--help", "-h", help=Messages.help),
)
def create(
    exists_ok: bool = False,
    no_start: bool = False,
    no_wait: bool = False,
    _show_help: bool = False,
    _extra: builtins.list[str] = [],
) -> UUID:
    """
    Create a new agent. The available agent recipes are fetched from your
    cluster and are added as dynamic subcommands. You can see more details
    and available arguments by calling the subcommand with --help, e.g. create
    [name] --help
    """
    auth = get_auth_state()
    settings = get_saved_settings()
    is_file_mode = settings.recipes_file is not None
    try:
        schemas = request_recipes(auth=auth, job_type="agent")
    except ProdigyTeamsError as e:
        raise CLIError(Messages.E009, e)
    except HTTPXErrors as e:
        raise CLIError(Messages.E009, e)
    if not _extra:
        print_recipes_help(schemas, "Create a new agent", "agents create")
        sys.exit(0)
    args = [*_extra]
    name = args.pop(0)
    if name not in schemas:
        opts = f"Available: {', '.join(schemas.keys())}"
        raise CLIError(Messages.E010.format(noun="agent", name=name), opts)
    if is_file_mode:
        schema = schemas[name]
        job_spec = create_from_recipe_file(
            schema, args, command="agents", show_help=_show_help, job_type="agent"
        )
        msg.good(Messages.T002.format(noun="agent", name=job_spec.name))
        print_args_table(job_spec.plan.args or {}, job_spec.plan.cli_names or {})
        if not no_start:
            response = auth.broker_client.jobs.start_job(job_spec)
            if response.cluster_change is not None:
                msg.good(Messages.T005.format(noun="agent", name=job_spec.job_id))
            elif response.cluster_error is not None:
                raise CLIError(
                    Messages.E025.format(noun="agent", name=job_spec.job_id),
                    response.cluster_error,
                )
        return job_spec.job_id
    schema = resolve_recipe(name, broker_id=auth.broker_id)
    agent_id, plan = create_from_recipe(
        schema, args, command="agents", show_help=_show_help, exists_ok=exists_ok
    )
    msg.good(Messages.T002.format(noun="agent", name=agent_id))
    print_args_table(plan.args, schema.form_schema.cli_names)
    if not no_start:
        global _pending_recipe_schema
        _pending_recipe_schema = schema
        start(agent_id, no_wait=no_wait)
        _pending_recipe_schema = None
    return agent_id


@cli.subcommand(
    "agents",
    "assign",
    name_or_id=Arg(help=Messages.name_or_id.format(noun="agent")),
    task=Arg("--task", help="Task name or ID to assign the agent to"),
    project_id=Arg(help=Messages.project_id.format(noun="agent")),
    cluster_id=Arg(help=Messages.cluster_id.format(noun="agent")),
)
def assign(
    name_or_id: Union[str, UUID],
    task: Union[str, UUID],
    project_id: Optional[UUID] = None,
    cluster_id: Optional[UUID] = None,
) -> None:
    """Assign an agent to a task as an annotator"""
    auth = get_auth_state()
    agent = resolve_agent(name_or_id, broker_id=cluster_id, project_id=project_id)
    task_obj = resolve_task(task, broker_id=cluster_id, project_id=project_id)
    auth.client.task.add_agent(task_id=task_obj.id, agent_id=agent.id)
    msg.good(f"Assigned agent '{agent.name}' to task '{task_obj.name}'")


@cli.subcommand(
    "agents",
    "list",
    # fmt: off
    select=Arg("--select", help=Messages.select.format(opts=[*builtins.list(AgentSummary.model_fields), "state"])),
    as_json=Arg("--json", help=Messages.as_json),
    # fmt: on
)
def list(
    select: builtins.list[str] = ["id", "name", "state", "project_name"],
    as_json: bool = False,
) -> Sequence[AgentSummary]:
    """
    List the agents on the cluster. By default, this includes their ID, name
    and current state, e.g. created or completed
    """
    auth = get_auth_state()
    res = auth.client.agent.all(page=1, size=100)
    items = res.items
    if "state" in select and items:
        job_ids = [item.id for item in items]
        try:
            statuses = auth.broker_client.jobs.get_batch_status(job_ids).statuses
        except (BrokerError, *HTTPXErrors):
            statuses = {}
        enriched = []
        for item in items:
            d = item.model_dump(include=set(s for s in select if s != "state"))
            status = statuses.get(str(item.id))
            d["state"] = status.state.value if status else "unknown"
            enriched.append(d)
        print_table_with_select(enriched, select=select, as_json=as_json)
    else:
        print_table_with_select(items, select=select, as_json=as_json)
    return items


@cli.subcommand(
    "agents",
    "info",
    name_or_id=Arg(help=Messages.name_or_id.format(noun="agent")),
    project_id=Arg(help=Messages.project_id.format(noun="agent")),
    cluster_id=Arg(help=Messages.cluster_id.format(noun="agent")),
    select=Arg(
        "--select",
        help=Messages.select.format(opts=builtins.list(AgentDetail.model_fields)),
    ),
    as_json=Arg("--json", help=Messages.as_json),
)
def info(
    name_or_id: Union[str, UUID],
    cluster_id: Optional[UUID] = None,
    project_id: Optional[UUID] = None,
    select: Optional[builtins.list[str]] = None,
    as_json: bool = False,
) -> AgentDetail:
    """Print information about an agent on the cluster"""
    auth = get_auth_state()
    res = resolve_agent(name_or_id, broker_id=cluster_id, project_id=project_id)
    try:
        status = auth.broker_client.jobs.get_status(res.id)
        msg.text(f"State: {status.state.value}", icon="info")
    except (BrokerError, *HTTPXErrors):
        pass
    print_info_table(res, exclude=["plan"], as_json=as_json, select=select)
    return res


@cli.subcommand(
    "agents",
    "logs",
    name_or_id=Arg(help=Messages.name_or_id_optional.format(noun="agent")),
    project_id=Arg(help=Messages.project_id.format(noun="agent")),
    cluster_id=Arg(help=Messages.cluster_id.format(noun="agent")),
    as_json=Arg("--json", help=Messages.as_json),
    show_errors=Arg(
        "--errors", help="Show structured error information instead of full logs"
    ),
    query=Arg("--query", help="Filter log lines matching this text (requires Loki)"),
)
def logs(
    name_or_id: Optional[Union[str, UUID]] = None,
    project_id: Optional[UUID] = None,
    cluster_id: Optional[UUID] = None,
    as_json: bool = False,
    show_errors: bool = False,
    query: Optional[str] = None,
) -> str | None:
    """Get logs for an agent on the cluster"""
    job = resolve_agent(name_or_id, broker_id=cluster_id, project_id=project_id)
    auth = get_auth_state()
    # Use last_execution_id if available (matches execution-id K8s label),
    # otherwise fall back to agent id (matches job-id K8s label).
    log_id = getattr(job, "last_execution_id", None) or job.id
    if show_errors:
        try:
            resp = auth.broker_client.jobs.errors(job_id=log_id)
        except BrokerError as e:
            raise CLIError(Messages.E011.format(noun="agent"), e)
        if as_json:
            print(resp.model_dump_json(indent=2))
        else:
            msg.text(f"State: {resp.state}", icon="info")
            if resp.errors:
                for err in resp.errors:
                    msg.fail(err.message)
                    if err.traceback:
                        print(err.traceback)
                        print()
            elif resp.tail:
                msg.warn("No structured errors found. Last 50 lines:")
                print(resp.tail)
            else:
                msg.good("No errors found")
        return None
    try:
        text = auth.broker_client.jobs.logs(job_id=log_id, query=query)
    except BrokerError as e:
        raise CLIError(Messages.E011.format(noun="agent"), e)
    print_logs(text, as_json=as_json)
    return text


@cli.subcommand(
    "agents",
    "start",
    name_or_id=Arg(help=Messages.name_or_id_optional.format(noun="agent")),
    project_id=Arg(help=Messages.project_id.format(noun="agent")),
    cluster_id=Arg(help=Messages.cluster_id.format(noun="agent")),
    worker_class=Arg(
        "--worker-class",
        help=Messages.recipe_worker_class.format(noun="agent"),
    ),
    no_wait=Arg("--no-wait", help=Messages.no_wait.format(noun="agent")),
)
def start(
    name_or_id: Optional[Union[str, UUID]] = None,
    project_id: Optional[UUID] = None,
    cluster_id: Optional[UUID] = None,
    worker_class: Optional[str] = None,
    no_wait: bool = False,
) -> UUID:
    """Start an agent on the cluster"""
    job = resolve_agent(name_or_id, broker_id=cluster_id, project_id=project_id)
    start_job(
        job,
        worker_class,
        get_auth_state(),
        recipe_schema=_pending_recipe_schema,
        wait=not no_wait,
    )
    return job.id


@cli.subcommand(
    "agents",
    "stop",
    name_or_id=Arg(help=Messages.name_or_id_optional.format(noun="agent")),
    project_id=Arg(help=Messages.project_id.format(noun="agent")),
    cluster_id=Arg(help=Messages.cluster_id.format(noun="agent")),
)
def stop(
    name_or_id: Optional[Union[str, UUID]] = None,
    project_id: Optional[UUID] = None,
    cluster_id: Optional[UUID] = None,
) -> UUID:
    """Stop an agent on the cluster"""
    job = resolve_agent(name_or_id, broker_id=cluster_id, project_id=project_id)
    stop_job(job, get_auth_state())
    return job.id


@cli.subcommand(
    "agents",
    "delete",
    name_or_id=Arg(help=Messages.name_or_id.format(noun="agent")),
    project_id=Arg(help=Messages.project_id.format(noun="agent")),
    cluster_id=Arg(help=Messages.cluster_id.format(noun="agent")),
)
def delete(
    name_or_id: Union[str, UUID],
    project_id: Optional[UUID] = None,
    cluster_id: Optional[UUID] = None,
) -> UUID:
    """Delete an agent by name or ID"""
    job = resolve_agent(name_or_id, broker_id=cluster_id, project_id=project_id)
    delete_job(job, get_auth_state())
    return job.id
