import json
import typing

import pydantic.json
import structlog

from skyvern.forge.sdk.artifact.models import Artifact, ArtifactType
from skyvern.forge.sdk.db.enums import OrganizationAuthTokenType
from skyvern.forge.sdk.db.models import (
    ArtifactModel,
    AWSSecretParameterModel,
    OrganizationAuthTokenModel,
    OrganizationModel,
    StepModel,
    TaskModel,
    WorkflowModel,
    WorkflowParameterModel,
    WorkflowRunModel,
    WorkflowRunParameterModel,
)
from skyvern.forge.sdk.models import Organization, OrganizationAuthToken, Step, StepStatus
from skyvern.forge.sdk.schemas.tasks import ProxyLocation, Task, TaskStatus
from skyvern.forge.sdk.workflow.models.parameter import AWSSecretParameter, WorkflowParameter, WorkflowParameterType
from skyvern.forge.sdk.workflow.models.workflow import (
    Workflow,
    WorkflowDefinition,
    WorkflowRun,
    WorkflowRunParameter,
    WorkflowRunStatus,
)

LOG = structlog.get_logger()


@typing.no_type_check
def _custom_json_serializer(*args, **kwargs) -> str:
    """
    Encodes json in the same way that pydantic does.
    """
    return json.dumps(*args, default=pydantic.json.pydantic_encoder, **kwargs)


def convert_to_task(task_obj: TaskModel, debug_enabled: bool = False) -> Task:
    if debug_enabled:
        LOG.debug("Converting TaskModel to Task", task_id=task_obj.task_id)
    task = Task(
        task_id=task_obj.task_id,
        status=TaskStatus(task_obj.status),
        created_at=task_obj.created_at,
        modified_at=task_obj.modified_at,
        title=task_obj.title,
        url=task_obj.url,
        webhook_callback_url=task_obj.webhook_callback_url,
        navigation_goal=task_obj.navigation_goal,
        data_extraction_goal=task_obj.data_extraction_goal,
        navigation_payload=task_obj.navigation_payload,
        extracted_information=task_obj.extracted_information,
        failure_reason=task_obj.failure_reason,
        organization_id=task_obj.organization_id,
        proxy_location=ProxyLocation(task_obj.proxy_location) if task_obj.proxy_location else None,
        extracted_information_schema=task_obj.extracted_information_schema,
        workflow_run_id=task_obj.workflow_run_id,
        order=task_obj.order,
        retry=task_obj.retry,
        error_code_mapping=task_obj.error_code_mapping,
        errors=task_obj.errors,
    )
    return task


def convert_to_step(step_model: StepModel, debug_enabled: bool = False) -> Step:
    if debug_enabled:
        LOG.debug("Converting StepModel to Step", step_id=step_model.step_id)
    return Step(
        task_id=step_model.task_id,
        step_id=step_model.step_id,
        created_at=step_model.created_at,
        modified_at=step_model.modified_at,
        status=StepStatus(step_model.status),
        output=step_model.output,
        order=step_model.order,
        is_last=step_model.is_last,
        retry_index=step_model.retry_index,
        organization_id=step_model.organization_id,
        input_token_count=step_model.input_token_count,
        output_token_count=step_model.output_token_count,
        step_cost=step_model.step_cost,
    )


def convert_to_organization(org_model: OrganizationModel) -> Organization:
    return Organization(
        organization_id=org_model.organization_id,
        organization_name=org_model.organization_name,
        webhook_callback_url=org_model.webhook_callback_url,
        max_steps_per_run=org_model.max_steps_per_run,
        created_at=org_model.created_at,
        modified_at=org_model.modified_at,
    )


def convert_to_organization_auth_token(org_auth_token: OrganizationAuthTokenModel) -> OrganizationAuthToken:
    return OrganizationAuthToken(
        id=org_auth_token.id,
        organization_id=org_auth_token.organization_id,
        token_type=OrganizationAuthTokenType(org_auth_token.token_type),
        token=org_auth_token.token,
        valid=org_auth_token.valid,
        created_at=org_auth_token.created_at,
        modified_at=org_auth_token.modified_at,
    )


def convert_to_artifact(artifact_model: ArtifactModel, debug_enabled: bool = False) -> Artifact:
    if debug_enabled:
        LOG.debug("Converting ArtifactModel to Artifact", artifact_id=artifact_model.artifact_id)

    return Artifact(
        artifact_id=artifact_model.artifact_id,
        artifact_type=ArtifactType[artifact_model.artifact_type.upper()],
        uri=artifact_model.uri,
        task_id=artifact_model.task_id,
        step_id=artifact_model.step_id,
        created_at=artifact_model.created_at,
        modified_at=artifact_model.modified_at,
        organization_id=artifact_model.organization_id,
    )


def convert_to_workflow(workflow_model: WorkflowModel, debug_enabled: bool = False) -> Workflow:
    if debug_enabled:
        LOG.debug("Converting WorkflowModel to Workflow", workflow_id=workflow_model.workflow_id)

    return Workflow(
        workflow_id=workflow_model.workflow_id,
        organization_id=workflow_model.organization_id,
        title=workflow_model.title,
        description=workflow_model.description,
        workflow_definition=WorkflowDefinition.model_validate(workflow_model.workflow_definition),
        created_at=workflow_model.created_at,
        modified_at=workflow_model.modified_at,
        deleted_at=workflow_model.deleted_at,
    )


def convert_to_workflow_run(workflow_run_model: WorkflowRunModel, debug_enabled: bool = False) -> WorkflowRun:
    if debug_enabled:
        LOG.debug("Converting WorkflowRunModel to WorkflowRun", workflow_run_id=workflow_run_model.workflow_run_id)

    return WorkflowRun(
        workflow_run_id=workflow_run_model.workflow_run_id,
        workflow_id=workflow_run_model.workflow_id,
        status=WorkflowRunStatus[workflow_run_model.status],
        proxy_location=ProxyLocation(workflow_run_model.proxy_location) if workflow_run_model.proxy_location else None,
        webhook_callback_url=workflow_run_model.webhook_callback_url,
        created_at=workflow_run_model.created_at,
        modified_at=workflow_run_model.modified_at,
    )


def convert_to_workflow_parameter(
    workflow_parameter_model: WorkflowParameterModel, debug_enabled: bool = False
) -> WorkflowParameter:
    if debug_enabled:
        LOG.debug(
            "Converting WorkflowParameterModel to WorkflowParameter",
            workflow_parameter_id=workflow_parameter_model.workflow_parameter_id,
        )

    workflow_parameter_type = WorkflowParameterType[workflow_parameter_model.workflow_parameter_type.upper()]

    return WorkflowParameter(
        workflow_parameter_id=workflow_parameter_model.workflow_parameter_id,
        workflow_parameter_type=workflow_parameter_type,
        workflow_id=workflow_parameter_model.workflow_id,
        default_value=workflow_parameter_type.convert_value(workflow_parameter_model.default_value),
        key=workflow_parameter_model.key,
        description=workflow_parameter_model.description,
        created_at=workflow_parameter_model.created_at,
        modified_at=workflow_parameter_model.modified_at,
        deleted_at=workflow_parameter_model.deleted_at,
    )


def convert_to_aws_secret_parameter(
    aws_secret_parameter_model: AWSSecretParameterModel, debug_enabled: bool = False
) -> AWSSecretParameter:
    if debug_enabled:
        LOG.debug(
            "Converting AWSSecretParameterModel to AWSSecretParameter",
            aws_secret_parameter_id=aws_secret_parameter_model.id,
        )

    return AWSSecretParameter(
        aws_secret_parameter_id=aws_secret_parameter_model.aws_secret_parameter_id,
        workflow_id=aws_secret_parameter_model.workflow_id,
        key=aws_secret_parameter_model.key,
        description=aws_secret_parameter_model.description,
        aws_key=aws_secret_parameter_model.aws_key,
        created_at=aws_secret_parameter_model.created_at,
        modified_at=aws_secret_parameter_model.modified_at,
        deleted_at=aws_secret_parameter_model.deleted_at,
    )


def convert_to_workflow_run_parameter(
    workflow_run_parameter_model: WorkflowRunParameterModel,
    workflow_parameter: WorkflowParameter,
    debug_enabled: bool = False,
) -> WorkflowRunParameter:
    if debug_enabled:
        LOG.debug(
            "Converting WorkflowRunParameterModel to WorkflowRunParameter",
            workflow_run_id=workflow_run_parameter_model.workflow_run_id,
            workflow_parameter_id=workflow_run_parameter_model.workflow_parameter_id,
        )

    return WorkflowRunParameter(
        workflow_run_id=workflow_run_parameter_model.workflow_run_id,
        workflow_parameter_id=workflow_run_parameter_model.workflow_parameter_id,
        value=workflow_parameter.workflow_parameter_type.convert_value(workflow_run_parameter_model.value),
        created_at=workflow_run_parameter_model.created_at,
    )