from __future__ import annotations

import time
from collections.abc import Callable, Iterable
from typing import Literal, TypeVar, Union
from uuid import UUID

import httpx
from pydantic import BaseModel
from wasabi import msg

from prodigy_teams_broker_sdk import Client as BrokerClient
from prodigy_teams_broker_sdk import models as broker_models
from prodigy_teams_broker_sdk.errors import BrokerError
from prodigy_teams_pam_sdk import Client as PamClient
from prodigy_teams_pam_sdk import models as pam_models
from prodigy_teams_pam_sdk.client.action import Action as PamActionEndpoint
from prodigy_teams_pam_sdk.client.agent import Agent as PamAgentEndpoint
from prodigy_teams_pam_sdk.client.task import Task as PamTaskEndpoint
from prodigy_teams_pam_sdk.errors import PRODIGY_TEAMS_ERRORS, ProdigyTeamsError

from .auth import AuthState
from .commands._state import get_auth_state, get_root_cfg, get_saved_settings
from .errors import CLIError, ProdigyTeamsErrors
from .messages import Messages
from .ty import Page
from .util import URL


def get_not_found_error(model_name: str) -> type[ProdigyTeamsError]:
    name = model_name.title().replace("_", "")
    return PRODIGY_TEAMS_ERRORS[f"{name}NotFound"]


def resolve_coding_task(
    id: UUID,
) -> pam_models.CodingTaskDetail:
    obj = _resolve_object("coding_task", id, {})
    assert isinstance(obj, pam_models.CodingTaskDetail)
    return obj


def resolve_action(
    name_or_id: str | UUID | None,
    *,
    broker_id: UUID | None,
    project_id: UUID | None,
) -> pam_models.ActionDetail:
    if name_or_id is None:
        settings = get_saved_settings()
        name_or_id = settings.action
        if name_or_id is None:
            raise ValueError("Action ID not set and no default")
    obj = _resolve_object(
        "action", name_or_id, {"broker_id": broker_id, "project_id": project_id}
    )
    assert isinstance(obj, pam_models.ActionDetail)
    # Fix URLs. It would be nice to do this in the actual models
    obj.url = URL.parse(obj.url).url
    obj.url_logs = URL.parse(obj.url_logs).url
    return obj


def resolve_action_id(
    name_or_id: str | UUID | None,
    *,
    broker_id: UUID | None,
    project_id: UUID | None,
) -> UUID:
    if name_or_id is None:
        settings = get_saved_settings()
        name_or_id = settings.action
        if name_or_id is None:
            raise ValueError("Action ID not set and no default")
    return _resolve_id(
        "action", name_or_id, {"broker_id": broker_id, "project_id": project_id}
    )


def resolve_agent(
    name_or_id: str | UUID | None,
    *,
    broker_id: UUID | None,
    project_id: UUID | None,
) -> pam_models.AgentDetail:
    if name_or_id is None:
        settings = get_saved_settings()
        name_or_id = settings.agent
        if name_or_id is None:
            raise ValueError("Agent ID not set and no default")
    obj = _resolve_object(
        "agent", name_or_id, {"broker_id": broker_id, "project_id": project_id}
    )
    assert isinstance(obj, pam_models.AgentDetail)
    return obj


def resolve_agent_id(
    name_or_id: str | UUID | None,
    *,
    broker_id: UUID | None,
    project_id: UUID | None,
) -> UUID:
    if name_or_id is None:
        settings = get_saved_settings()
        name_or_id = settings.agent
        if name_or_id is None:
            raise ValueError("Agent ID not set and no default")
    return _resolve_id(
        "agent", name_or_id, {"broker_id": broker_id, "project_id": project_id}
    )


def resolve_asset(
    name_or_id: str | UUID,
    *,
    broker_id: UUID | None,
    project_id: UUID | None,
) -> pam_models.AssetDetail:
    obj = _resolve_object(
        "asset", name_or_id, {"broker_id": broker_id, "project_id": project_id}
    )
    assert isinstance(obj, pam_models.AssetDetail)
    return obj


def resolve_asset_id(
    name_or_id: str | UUID,
    *,
    broker_id: UUID | None,
    project_id: UUID | None,
) -> UUID:
    return _resolve_id(
        "asset", name_or_id, {"broker_id": broker_id, "project_id": project_id}
    )


def resolve_broker(
    name_or_id: str | UUID,
) -> pam_models.BrokerDetail:
    obj = _resolve_object("broker", name_or_id, {})
    assert isinstance(obj, pam_models.BrokerDetail)
    return obj


def resolve_broker_id(
    name_or_id: str | UUID | None,
) -> UUID:
    if name_or_id is None:
        settings = get_saved_settings()
        name_or_id = settings.broker_id
        if name_or_id is None:
            raise ValueError("Cluster ID not set and no default")
    return _resolve_id("broker", name_or_id, {})


def resolve_dataset(
    name_or_id: str | UUID,
    *,
    broker_id: UUID | None,
) -> pam_models.DatasetDetail:
    obj = _resolve_object("dataset", name_or_id, {"broker_id": broker_id})
    assert isinstance(obj, pam_models.DatasetDetail)
    return obj


def resolve_dataset_id(
    name_or_id: str | UUID,
    *,
    broker_id: UUID | None,
) -> UUID:
    return _resolve_id("dataset", name_or_id, {"broker_id": broker_id})


def resolve_package(
    name_or_id: str | UUID,
    *,
    broker_id: UUID | None,
) -> pam_models.PackageDetail:
    obj = _resolve_latest_object(
        "package",
        name_or_id,
        {"broker_id": broker_id},
        pam_models.PackageReadingLatest,
    )
    assert isinstance(obj, pam_models.PackageDetail)
    return obj


def resolve_package_id(
    name_or_id: str | UUID,
    *,
    broker_id: UUID | None,
) -> UUID:
    return _resolve_id("package", name_or_id, {"broker_id": broker_id})


def resolve_path(
    name_or_id: str | UUID,
    *,
    broker_id: UUID | None,
) -> pam_models.BrokerPathDetail:
    obj = _resolve_object("broker_path", name_or_id, {"broker_id": broker_id})
    assert isinstance(obj, pam_models.BrokerPathDetail)
    return obj


def resolve_path_id(
    name_or_id: str | UUID,
    *,
    broker_id: UUID | None,
) -> UUID:
    return _resolve_id("broker_path", name_or_id, {"broker_id": broker_id})


def resolve_project(
    name_or_id: str | UUID,
) -> pam_models.ProjectDetail:
    obj = _resolve_object("project", name_or_id, {})
    assert isinstance(obj, pam_models.ProjectDetail)
    return obj


def resolve_project_id(
    name_or_id: str | UUID,
) -> UUID:
    return _resolve_id("project", name_or_id, {})


def resolve_recipe(
    name_or_id: str | UUID,
    *,
    broker_id: UUID | None,
) -> pam_models.RecipeDetail:
    obj = _resolve_latest_object(
        "recipe",
        name_or_id,
        {"broker_id": broker_id},
        pam_models.RecipeReadingLatest,
    )
    assert isinstance(obj, pam_models.RecipeDetail)
    return obj


# TODO: this is currently not used
def resolve_recipe_id(
    name_or_id: str | UUID,
    *,
    broker_id: UUID | None,
) -> UUID:
    return _resolve_id("recipe", name_or_id, {"broker_id": broker_id})


def resolve_secret(
    name_or_id: str | UUID,
    *,
    broker_id: UUID | None,
) -> pam_models.SecretDetail:
    obj = _resolve_object("secret", name_or_id, {"broker_id": broker_id})
    assert isinstance(obj, pam_models.SecretDetail)
    return obj


def resolve_secret_id(
    name_or_id: str | UUID,
    *,
    broker_id: UUID | None,
) -> UUID:
    return _resolve_id("secret", name_or_id, {"broker_id": broker_id})


def resolve_task(
    name_or_id: str | UUID | None,
    *,
    broker_id: UUID | None,
    project_id: UUID | None,
) -> pam_models.TaskDetail:
    if name_or_id is None:
        settings = get_saved_settings()
        name_or_id = settings.task
        if name_or_id is None:
            raise ValueError("Task ID not set and no default")
    obj = _resolve_object(
        "task", name_or_id, {"broker_id": broker_id, "project_id": project_id}
    )
    assert isinstance(obj, pam_models.TaskDetail)
    # Fix URLs. It would be nice to do this in the actual models
    obj.url = URL.parse(obj.url).url
    obj.url_logs = URL.parse(obj.url_logs).url
    return obj


def resolve_task_id(
    name_or_id: str | UUID | None,
    *,
    broker_id: UUID | None,
    project_id: UUID | None,
) -> UUID:
    if name_or_id is None:
        settings = get_saved_settings()
        name_or_id = settings.task
        if name_or_id is None:
            raise ValueError("Task ID not set and no default")
    return _resolve_id(
        "task", name_or_id, {"broker_id": broker_id, "project_id": project_id}
    )


def _resolve_id(
    model_name: str,
    name_or_id: str | UUID,
    params: dict[str, str | UUID | None],
) -> UUID:
    """Resolve reference to a Prodigy Teams Entity by name or ID.
    If an ID is given, return it. If it's a name,
    look up the ID.
    """
    if isinstance(name_or_id, UUID):
        return name_or_id
    try:
        return UUID(name_or_id)
    except ValueError:
        pass
    auth = get_auth_state()
    model_client = getattr(auth.client, model_name)
    NotFound = get_not_found_error(model_name)
    try:
        res = model_client.read(name=name_or_id, **params)
        return res.id
    except NotFound:
        err = Messages.E038.format(noun=model_name, name_or_id=name_or_id)
        raise CLIError(err, params)


def _resolve_object(
    model_name: str,
    name_or_id: str | UUID,
    params: dict[str, str | UUID | None],
) -> BaseModel:
    """Resolve reference to a Prodigy Teams Entity by name or ID.
    If an ID is given, use it for an id query, otherwise, query by
    the entity name
    """
    auth = get_auth_state()
    model_client = getattr(auth.client, model_name)
    NotFound = get_not_found_error(model_name)
    record_id = None
    if isinstance(name_or_id, UUID):
        record_id = name_or_id
    else:
        try:
            record_id = UUID(name_or_id)
        except ValueError:
            pass
    if record_id:
        try:
            res = model_client.read(id=record_id, **params)
        except NotFound:
            err = Messages.E038.format(noun=model_name, name_or_id=name_or_id)
            raise CLIError(err, params)
    else:
        try:
            res = model_client.read(name=name_or_id, **params)
        except NotFound:
            err = Messages.E038.format(noun=model_name, name_or_id=name_or_id)
            raise CLIError(err, params)
    return res


def _resolve_latest_object(
    model_name: str,
    name_or_id: str | UUID,
    params: dict[str, str | UUID | None],
    latest_request_model: type[pam_models.BaseModel],
) -> pam_models.BaseModel:
    """Resolve latest reference to a Prodigy Teams Entity by name or ID.
    If an ID is given, use it for an id query, otherwise, query by
    the entity name
    """
    auth = get_auth_state()
    model_client = getattr(auth.client, model_name)
    NotFound = get_not_found_error(model_name)
    record_id = None
    if isinstance(name_or_id, UUID):
        record_id = name_or_id
    else:
        try:
            record_id = UUID(name_or_id)
        except ValueError:
            pass
    if record_id:
        try:
            res = model_client.read(id=record_id, **params)
        except NotFound:
            err = Messages.E038.format(noun=model_name, name_or_id=name_or_id)
            raise CLIError(err, params)
    else:
        body = latest_request_model(name=name_or_id, **params)
        try:
            res = model_client.latest(body=body)
        except NotFound:
            err = Messages.E038.format(
                noun=model_name, name_or_id=getattr(body, "name", None)
            )
            raise CLIError(err, body)
    return res


class JobOperations:
    pam_client: PamClient
    broker_client: BrokerClient
    job_type: broker_models.JobType

    def __init__(
        self,
        pam_client: PamClient,
        broker_client: BrokerClient,
        job_type: Literal["task", "action", "agent"],
    ) -> None:
        self.pam_client = pam_client
        self.broker_client = broker_client
        self.job_type = broker_models.JobType(job_type)

    def start(self, job_spec: broker_models.JobSpec, quiet: bool = False) -> None:
        change_response = self.broker_client.jobs.start_job(job_spec)
        if change_response.cluster_change is not None:
            return
        elif change_response.cluster_error:
            err = Messages.E025.format(noun=self.job_type.value, name=job_spec.job_id)
            raise CLIError(err, change_response.cluster_error)
        elif change_response.validation_error:
            err = Messages.E044.format(noun=self.job_type.value, name=job_spec.job_id)
            raise CLIError(err, change_response.validation_error)
        elif not quiet:
            msg.info(
                Messages.T006.format(noun=self.job_type.value, name=job_spec.job_id)
            )

    def stop(self, job_id: UUID, quiet: bool = False) -> None:
        change_response = self.broker_client.jobs.stop_job(
            broker_models.JobID(id=job_id, job_type=self.job_type)
        )
        if change_response.cluster_change is not None:
            if not quiet:
                msg.good(Messages.T007.format(noun=self.job_type.value, name=job_id))
        elif change_response.cluster_error:
            err = Messages.E026.format(noun=self.job_type.value, name=job_id)
            raise CLIError(err, change_response.cluster_error)
        elif change_response.validation_error:
            err = Messages.E044.format(noun=self.job_type.value, name=job_id)
            raise CLIError(err, change_response.validation_error)
        elif not quiet:
            msg.info(Messages.T008.format(noun=self.job_type.value, name=job_id))

    def delete(self, job_id: UUID, quiet: bool = False) -> None:
        try:
            change_response = self.broker_client.jobs.stop_job(
                broker_models.JobID(id=job_id, job_type=self.job_type)
            )
            if change_response.cluster_error:
                err = Messages.E028.format(noun=self.job_type.value, name=job_id)
                raise CLIError(err, change_response.cluster_error)
        except httpx.HTTPStatusError as e:
            if e.response.status_code == 500:
                # Backwards compatibility: older brokers can return 500 when the job has been GC'd
                pass
            else:
                raise
        else:
            if not quiet:
                msg.good(Messages.T009.format(noun=self.job_type.value, name=job_id))
        try:
            self.pam_endpoint().delete(id=job_id)
        except (
            ProdigyTeamsErrors.TaskNotFound,
            ProdigyTeamsErrors.ActionNotFound,
            ProdigyTeamsErrors.AgentNotFound,
        ):
            raise CLIError(Messages.E029.format(noun=self.job_type.value, name=job_id))
        except (
            ProdigyTeamsErrors.TaskForbiddenDelete,
            ProdigyTeamsErrors.TaskForbiddenRead,
            ProdigyTeamsErrors.ActionForbiddenDelete,
            ProdigyTeamsErrors.ActionForbiddenRead,
            ProdigyTeamsErrors.AgentForbiddenDelete,
            ProdigyTeamsErrors.AgentForbiddenRead,
        ):
            raise CLIError(Messages.E030.format(noun=self.job_type.value, name=job_id))
        if not quiet:
            msg.good(Messages.T003.format(noun=self.job_type.value, name=job_id))

    def pam_endpoint(self) -> PamActionEndpoint | PamAgentEndpoint | PamTaskEndpoint:
        return getattr(self.pam_client, self.job_type.value)


_JobT = TypeVar(
    "_JobT",
    bound=Union[pam_models.TaskDetail, pam_models.ActionDetail, pam_models.AgentDetail],
)


def _build_job_spec(
    job: pam_models.TaskDetail | pam_models.ActionDetail | pam_models.AgentDetail,
) -> broker_models.JobSpec:
    plan = job.plan
    if isinstance(job, pam_models.TaskDetail):
        job_type = "task"
    elif isinstance(job, pam_models.AgentDetail):
        job_type = "agent"
    else:
        job_type = "action"
    worker_type = getattr(plan, "worker_class", None)
    return broker_models.JobSpec(
        job_id=job.id,
        job_type=broker_models.JobType(job_type),
        name=job.name,
        plan=broker_models.RecipePlan(
            recipe_name=plan.recipe_name,
            args=plan.args,
            positional_fields=plan.positional_fields,
            cli_names=plan.cli_names,
            package_name=plan.package_name,
            objects=plan.objects,
        ),
        container=plan.container,
        node_selector=None,
        worker_type=worker_type,
    )


def _get_job_type(
    job: pam_models.TaskDetail | pam_models.ActionDetail | pam_models.AgentDetail,
) -> Literal["task", "action", "agent"]:
    if isinstance(job, pam_models.TaskDetail):
        return "task"
    elif isinstance(job, pam_models.AgentDetail):
        return "agent"
    return "action"


def start_job(
    job: _JobT,
    worker_class: str | None,
    auth: AuthState,
    wait: bool = True,
    quiet: bool = False,
) -> _JobT:
    job_type = _get_job_type(job)
    job_operations = JobOperations(auth.client, auth.broker_client, job_type)
    if worker_class is not None:
        job.plan = auth.client.recipeplan.update(
            pam_models.RecipePlanUpdating(id=job.plan.id, worker_class=worker_class)
        )
    job_spec = _build_job_spec(job)
    job_operations.start(job_spec, quiet=quiet)
    root_cfg = get_root_cfg()
    settings = get_saved_settings()
    settings.update(job_type, job.id)
    if wait:
        res = check_job_started(job, auth)
        if res is not None and isinstance(
            res, (pam_models.TaskDetail, pam_models.ActionDetail)
        ):
            res.url = URL.parse(res.url).url  # type: ignore[union-attr]
            res.url_logs = URL.parse(res.url_logs).url  # type: ignore[union-attr]
            msg.info(f"Starting {res.url}")
    # FIXME: this shouldn't overwrite any other settings
    settings.save(root_cfg.saved_settings_path)
    return job


def stop_job(job: _JobT, auth: AuthState, quiet: bool = False) -> _JobT:
    job_type = _get_job_type(job)
    job_operations = JobOperations(auth.client, auth.broker_client, job_type)
    job_operations.stop(job.id, quiet=quiet)
    root_cfg = get_root_cfg()
    settings = get_saved_settings()
    settings.update(job_type, job.id)
    # FIXME: this shouldn't overwrite any other settings
    settings.save(root_cfg.saved_settings_path)
    return job


def delete_job(job: _JobT, auth: AuthState, quiet: bool = False) -> _JobT:
    root_cfg = get_root_cfg()
    job_type = _get_job_type(job)
    job_operations = JobOperations(auth.client, auth.broker_client, job_type)
    job_operations.delete(job.id, quiet=quiet)
    root_cfg = get_root_cfg()
    settings = get_saved_settings()
    settings.update(job_type, None)
    # FIXME: this shouldn't overwrite any other settings
    settings.save(root_cfg.saved_settings_path)
    return job


def check_job_started(
    job: _JobT, auth: AuthState, *, wait_seconds: int = 60 * 2
) -> _JobT | None:
    job_type = _get_job_type(job)
    started_states = {
        broker_models.WorkloadState.running,
        broker_models.WorkloadState.ready,
        broker_models.WorkloadState.succeeded,
    }
    failed_states = {
        broker_models.WorkloadState.failed,
        broker_models.WorkloadState.not_found,
    }
    failed_healths = {
        broker_models.WorkloadHealth.degraded,
        broker_models.WorkloadHealth.terminal,
    }
    status: broker_models.JobStatusResponse | None = None
    from .ui import loading

    with loading(Messages.T034.format(noun=job_type)):
        for _ in range(wait_seconds):
            try:
                status = auth.broker_client.jobs.get_status(job.id)
            except (httpx.HTTPError, BrokerError):
                time.sleep(1)
                continue
            assert status is not None
            if status.is_ready or (
                status.state in started_states
                and getattr(status, "health", None) not in failed_healths
            ):
                break
            if status.state in failed_states or (
                getattr(status, "health", None) in failed_healths
            ):
                break
            time.sleep(1)
    if status is not None and (
        status.is_ready
        or (
            status.state in started_states
            and getattr(status, "health", None) not in failed_healths
        )
    ):
        msg.good(Messages.T005.format(noun=job_type, name=job.id))
        return job
    elif status is not None and (
        status.state in failed_states
        or getattr(status, "health", None) in failed_healths
    ):
        detail = status.message or getattr(job, "url_logs", None)
        raise CLIError(Messages.E025.format(noun=job_type, name=job.id), detail)
    msg.warn(
        Messages.E025.format(noun=job_type, name=job.id),
        getattr(job, "url_logs", None),
    )


def collect_from_pages(request: Callable[[int], Page]) -> Iterable[Page]:
    page = 1
    while True:
        page_result = request(page)

        if page_result.items:
            yield page_result
        else:
            break

        page += 1
