from __future__ import annotations

from collections.abc import Callable
from typing import cast
from uuid import UUID

import uuid_utils
from radicli import (
    DEFAULT_PLACEHOLDER,
    Arg,
    ArgparseArg,
    ArgumentParser,
    ConverterType,
    get_arg,
)
from wasabi import msg

from prodigy_teams_broker_sdk import models as broker_models
from prodigy_teams_broker_sdk.models import JobID, JobType
from prodigy_teams_pam_sdk import Client
from prodigy_teams_pam_sdk.models import (
    ActionCreating,
    AgentCreating,
    ObjectValidation,
    RecipeDetail,
    RecipeListingLatest,
    RecipePlanCreating,
    TaskCreating,
)
from prodigy_teams_pam_sdk.recipe_utils import (
    Field,
    RecipeSchema,
    get_field_cli_args,
    get_field_optional,
    parse_plan,
)

from ..auth import AuthState
from ..cli import CONVERTERS, cli
from ..config import SavedSettings
from ..errors import CLIError, ProdigyTeamsErrors
from ..messages import Messages
from ..query import resolve_project_id
from ..ui import RecipeHelpFormatter, get_recipe_desc
from ._recipe_file import load_recipes_file
from ._state import get_auth_state, get_root_cfg, get_saved_settings


def get_converter(arg_type: type) -> ConverterType | None:
    return CONVERTERS.get(arg_type, None)


def _get_noun(recipe: RecipeDetail) -> str:
    """Derive a human-readable noun from a recipe's job_type or is_action flag."""
    jt = getattr(recipe, "job_type", None)
    if jt in ("task", "action", "agent"):
        return jt
    return "action" if recipe.is_action else "task"


def request_recipes(
    auth: AuthState,
    is_action: bool | None = None,
    job_type: str | None = None,
) -> dict[str, RecipeDetail]:
    """Fetch all Recipes from Prodigy Teams or a local recipes file.
    auth (AuthState): CLI Auth object to handle auth to the Prodigy Teams API
    RETURNS (Dict[str, RecipeDetail]): Recipe schemas keyed by name
    """
    from pathlib import Path

    settings = get_saved_settings()
    if settings.recipes_file is not None:
        return load_recipes_file(Path(settings.recipes_file), is_action, job_type)
    client = auth.client
    if client is None:
        return {}
    from ..ui import loading

    with loading(Messages.T020):
        recipes = auth.retry(
            lambda: list(
                client.recipe.all_latest(
                    body=RecipeListingLatest(
                        org_id=auth.org_id,
                        broker_id=auth.broker_id,
                    )
                )
            )
        )()
    schemas = {}
    for recipe in recipes:
        if job_type is not None:
            if getattr(recipe, "job_type", None) == job_type:
                schemas[recipe.name] = recipe
        elif is_action is None or recipe.is_action == is_action:
            schemas[recipe.name] = recipe
    return schemas


def create_from_recipe(
    recipe: RecipeDetail,
    args: list[str],
    *,
    command: str,
    show_help: bool,
    exists_ok: bool = False,
) -> tuple[UUID, RecipePlanCreating]:
    # There's a few steps here:
    # 1. Parse the CLI string into arguments
    # 2. Create the record in the PAM DB
    # And then this is handled elsewhere:
    # 3. Tell the cluster to start the job
    try:
        task_or_action_name, project_id, plan = _parse_plan(
            recipe,
            args,
            parser_prog=f"{cli.prog} {command} create {recipe.name.lower()}",
            show_help=show_help,
            get_project_id=_get_default_project_id,
        )
    except ShowHelp as e:
        print(e.help_text)
        raise SystemExit(0)
    auth = get_auth_state()
    object_validation = _get_object_validation(plan)
    if object_validation.errors:
        errors = ["\n".join(err) for err in object_validation.errors.values()]
        noun = _get_noun(recipe)
        err = Messages.E044.format(noun=noun, name=task_or_action_name)
        raise CLIError(err, "\n".join(errors))
    _validate_plan(recipe, plan, object_validation)
    jt = getattr(recipe, "job_type", None)
    if jt == "agent":
        job_id = _create_agent(
            auth.client, task_or_action_name, project_id, plan, exists_ok=exists_ok
        )
    elif recipe.is_action:
        job_id = _create_action(
            auth.client, task_or_action_name, project_id, plan, exists_ok=exists_ok
        )
    else:
        job_id = _create_task(
            auth.client, task_or_action_name, project_id, plan, exists_ok=exists_ok
        )
    return job_id.id, plan


def _parse_plan(
    recipe: RecipeDetail,
    args: list[str],
    *,
    parser_prog: str,
    show_help: bool,
    get_project_id: Callable[[], UUID | None],
) -> tuple[str | None, UUID, RecipePlanCreating]:
    # This function works hard to be pure and not rely on
    # external state. This way we'll be able to work with it in
    # the compatibility tests more easily.
    noun = _get_noun(recipe)
    ap_args = _get_argparse_args(recipe)
    parser = _get_argparse_parser(parser_prog, recipe, ap_args)
    if not args or show_help:
        raise ShowHelp(parser.format_help())
    task_or_action_name, project_id, worker_class, plan_args = _parse_cli_args(
        parser, args
    )
    missing = _check_missing(recipe, ap_args, plan_args)
    if missing:
        err = Messages.E001.format(noun=noun, name=task_or_action_name)
        raise CLIError(err, Messages.E043.format(args=", ".join(missing)))
    if project_id is None:
        project_id = get_project_id()
        if project_id is None:
            raise CLIError(Messages.E042.format(noun=noun, name=task_or_action_name))

    plan = RecipePlanCreating(
        recipe_id=recipe.id,
        project_id=project_id,
        args=plan_args,
        worker_class=worker_class,
    )
    return task_or_action_name, project_id, plan


def _get_default_project_id() -> UUID | None:
    root_cfg = get_root_cfg()
    project_id = SavedSettings.from_file(root_cfg.saved_settings_path).project
    if project_id is not None:
        project_id = resolve_project_id(project_id)
    return project_id


def _parse_cli_args(
    parser: ArgumentParser,
    args: list[str],
) -> tuple[str | None, UUID | None, str | None, dict[str, Arg]]:
    parsed = vars(parser.parse_args(args))
    task_or_action_name = parsed.pop("name", None)
    assert isinstance(task_or_action_name, (str, type(None)))
    project_id = parsed.pop("project_id")
    if project_id is not None:
        project_id = UUID(project_id)
    worker_class = parsed.pop("worker_class", None)
    parsed.pop("runtime", None)  # legacy field, ignored
    return (task_or_action_name, project_id, worker_class, parsed)


def _check_missing(recipe: RecipeDetail, ap_args, parsed) -> list[str]:
    optional = get_field_optional(recipe.form_schema)  # type: ignore[arg-type]
    missing = []
    # We're iterating over the original args here because that gives us access
    # to the option names in the right CLI syntax
    for arg in ap_args:
        required = not optional.get(arg.id, True)
        if required and arg.id not in parsed:
            missing.append(arg.arg.option or arg.id)
    return missing


def _get_argparse_args(recipe: RecipeDetail) -> list[ArgparseArg]:
    noun = _get_noun(recipe)
    ap_args: list[ArgparseArg] = []
    ap_args.append(
        get_arg(
            "name",
            Arg("--name", help=Messages.name.format(noun=noun)),
            str,
            default=None,
        )
    )
    for field in recipe.form_schema.fields:
        ap_args.extend(
            get_field_cli_args(cast(Field, field), recipe.form_schema.cli_names or {})
        )
    return ap_args


def _get_argparse_parser(
    prog: str, recipe: RecipeDetail, ap_args: list[ArgparseArg]
) -> ArgumentParser:
    noun = _get_noun(recipe)
    # To get argument parsing and consistent help docs out-of-the-box we create
    # a dummy parser with radicli's modified ArgumentParser (argparse subclass)
    p = ArgumentParser(
        prog=prog,
        description=get_recipe_desc(recipe),
        formatter_class=RecipeHelpFormatter,
    )
    for arg in ap_args:
        a_args, a_kwargs = arg.to_argparse()
        # We don't want defaults to creep in here, we want that to be
        # introduced at the last moment by the actual function.
        if "default" in a_kwargs:
            a_kwargs["default"] = DEFAULT_PLACEHOLDER
        p.add_argument(*a_args, **a_kwargs)
    # Add default arguments to all recipe CLIs
    pid_arg = Arg("--project-id", help=Messages.recipe_project_id.format(noun=noun))
    pid = get_arg("project_id", pid_arg, str, default=None)
    pid_args, pid_kwargs = pid.to_argparse()
    p.add_argument(*pid_args, **pid_kwargs)
    worker_class_arg = Arg(
        "--worker-class",
        help=Messages.recipe_worker_class.format(noun=noun),
    )
    worker_class = get_arg("worker_class", worker_class_arg, str, default=None)
    worker_class_args, worker_class_kwargs = worker_class.to_argparse()
    p.add_argument(*worker_class_args, **worker_class_kwargs)
    runtime_arg = Arg(
        "--runtime",
        help=Messages.recipe_runtime.format(noun=noun),
    )
    runtime = get_arg("runtime", runtime_arg, str, default=None)
    runtime_args, runtime_kwargs = runtime.to_argparse()
    p.add_argument(*runtime_args, **runtime_kwargs)
    return p


def _get_object_validation(plan: RecipePlanCreating) -> ObjectValidation:
    auth = get_auth_state()
    object_validation = auth.client.recipeplan.validate_objects(plan)
    return object_validation


def _validate_plan(
    recipe: RecipeDetail,
    plan: RecipePlanCreating,
    object_validation: ObjectValidation,
) -> None:
    # Building a dummy objects map for parse_plan – the values don't matter,
    # since we're also only using a dummy fuction for create_custom_type
    objects_map = {"asset": {}, "dataset": {}, "secret": {}}
    for field_type, objects in object_validation.existing.items():
        objects_map[field_type] = {arg: {} for arg in objects}
    for field_type, objects in object_validation.to_create.items():
        objects_map.setdefault(field_type, {}).update({arg: {} for arg in objects})
    parse_plan(
        cast(RecipeSchema, recipe.form_schema),
        plan.args,
        objects_map,
        lambda x: lambda *a, **b: x,
    )


def _create_task(
    client: Client,
    name: str | None,
    project_id: UUID,
    plan: RecipePlanCreating,
    exists_ok: bool = False,
) -> JobID:
    body = TaskCreating(name=name, project_id=project_id, plan=plan)
    try:
        task = client.task.create(body)
    except ProdigyTeamsErrors.ProjectNotFound:
        err = Messages.E013.format(name=project_id)
        raise CLIError(err)
    except ProdigyTeamsErrors.TaskExists:
        if exists_ok:
            msg.info(Messages.T001.format(noun="task", name=name))
            task = client.task.read(name=name, project_id=project_id)
            return JobID(id=task.id, job_type=JobType.task)
        err = Messages.E014.format(noun="task", name=name, project=project_id)
        raise CLIError(err)
    except ProdigyTeamsErrors.TaskInvalid:
        raise CLIError(Messages.E004.format(noun="task", name=name))
    return JobID(id=task.id, job_type=JobType.task)


def _create_action(
    client: Client,
    name: str | None,
    project_id: UUID,
    plan: RecipePlanCreating,
    exists_ok: bool = False,
) -> JobID:
    body = ActionCreating(name=name, project_id=project_id, evaluation="", plan=plan)
    try:
        action = client.action.create(body)
    except ProdigyTeamsErrors.ProjectNotFound:
        err = Messages.E013.format(name=project_id)
        raise CLIError(err)
    except ProdigyTeamsErrors.ActionExists:
        if exists_ok:
            msg.info(Messages.T001.format(noun="action", name=name))
            action = client.action.read(name=name, project_id=project_id)
            return JobID(id=action.id, job_type=JobType.action)
        err = Messages.E014.format(noun="action", name=name, project=project_id)
        raise CLIError(err)
    except ProdigyTeamsErrors.ActionInvalid:
        raise CLIError(Messages.E004.format(noun="action", name=name))
    return JobID(id=action.id, job_type=JobType.action)


def _create_agent(
    client: Client,
    name: str | None,
    project_id: UUID,
    plan: RecipePlanCreating,
    exists_ok: bool = False,
) -> JobID:
    body = AgentCreating(name=name, project_id=project_id, plan=plan)
    try:
        agent = client.agent.create(body)
    except ProdigyTeamsErrors.ProjectNotFound:
        err = Messages.E013.format(name=project_id)
        raise CLIError(err)
    except ProdigyTeamsErrors.AgentExists:
        if exists_ok:
            msg.info(Messages.T001.format(noun="agent", name=name))
            agent = client.agent.read(name=name, project_id=project_id)
            return JobID(id=agent.id, job_type=JobType.agent)
        err = Messages.E014.format(noun="agent", name=name, project=project_id)
        raise CLIError(err)
    except ProdigyTeamsErrors.AgentInvalid:
        raise CLIError(Messages.E004.format(noun="agent", name=name))
    return JobID(id=agent.id, job_type=JobType.agent)


def _uuid7() -> UUID:
    """Generate a UUIDv7 (time-sortable) as a stdlib uuid.UUID."""
    return UUID(int=uuid_utils.uuid7().int)


def create_from_recipe_file(
    recipe: RecipeDetail,
    args: list[str],
    *,
    command: str,
    show_help: bool,
    job_type: str,
) -> broker_models.JobSpec:
    """Parse CLI args and build a JobSpec directly, bypassing PAM.

    Used when a local recipes file is configured.
    """
    ap_args = _get_argparse_args(recipe)
    parser = _get_argparse_parser(
        f"{cli.prog} {command} create {recipe.name.lower()}", recipe, ap_args
    )
    if not args or show_help:
        print(parser.format_help())
        raise SystemExit(0)
    parsed = vars(parser.parse_args(args))
    job_name = parsed.pop("name", None)
    parsed.pop("project_id", None)
    parsed.pop("worker_class", None)
    # Remove DEFAULT_PLACEHOLDER values (unfilled optional args)
    plan_args = {k: v for k, v in parsed.items() if v is not DEFAULT_PLACEHOLDER}
    # Extract positional field names
    positional_fields = [
        f.name for f in recipe.form_schema.fields if f.props.positional
    ]
    cli_names = recipe.form_schema.cli_names or {}
    if not recipe.package or not recipe.package.container:
        raise CLIError("Recipe has no associated package container")
    recipe_plan = broker_models.RecipePlan(
        recipe_name=recipe.name,
        args=plan_args,
        positional_fields=positional_fields,
        cli_names=cli_names,
        package_name=recipe.package.name,
        objects={},
    )
    if job_name is None:
        job_name = recipe.name
    return broker_models.JobSpec(
        job_id=_uuid7(),
        job_type=broker_models.JobType(job_type),
        name=job_name,
        plan=recipe_plan,
        container=recipe.package.container,
        resources=None,
        node_selector=None,
    )


class ShowHelp(ValueError):
    def __init__(self, help_text: str):
        self.help_text = help_text
