from __future__ import annotations

import functools
import logging
import os
from contextlib import nullcontext
from copy import deepcopy
from dataclasses import replace
from typing import Any, TypedDict

from agente_contas_tim.observability import (
    emit_workflow_step,
    langfuse_context_has_active_span,
)
from agente_contas_tim.workflows.actions.registry import (
    ActionRegistry,
    WorkflowRuntimeContext,
)
from agente_contas_tim.workflows.conditions import (
    evaluate_condition,
    resolve_value,
)
from agente_contas_tim.workflows.contracts import (
    EdgeDef,
    NodeDef,
    PauseDef,
    WorkflowDef,
)
from agente_contas_tim.workflows.exceptions import (
    WorkflowConfigurationError,
    WorkflowInputError,
)
from agente_contas_tim.workflows.runtime_types import ActionResult
from agente_contas_tim.workflows.templating import render_template

logger = logging.getLogger(__name__)

FINISH_NODE = "__workflow_finish__"

_ACTION_LABELS: dict[str, str] = {
    "consulta_vas": "Consulta VAS na linha",
    "cancelar_vas_single": "Cancelamento do serviço",
    "consultar_perfil_fatura": "Consulta tipo de fatura",
    "check_invoice_status": "Check invoice status",
    "enviar_sms": "Envio de SMS com boleto",
    "atualizar_status_sr": "Atualização status SR",
    "consultar_divergencia": "Consulta divergência",
    "bloquear_vas": "Bloqueio de VAS",
    "bloquear_vas_single": "Bloqueio de VAS",
    "cancelar_vas": "Cancelamento de VAS",
    "avaliar_proxima_acao": "Avaliação próxima ação",
    "resolve_capability": "Resolução de capability",
    "preparar_vas_estrategico": "Preparacao VAS estrategico",
    "montar_explicacao_cancelamento_vas_estrategico": (
        "Explicacao cancelamento VAS estrategico"
    ),
    "registrar_atendimento_vas_estrategico": (
        "Registro de atendimento VAS estrategico"
    ),
}


@functools.lru_cache(maxsize=1)
def _import_langfuse_get_client() -> Any:
    from langfuse import get_client
    return get_client


@functools.lru_cache(maxsize=1)
def _import_langfuse_callback_handler() -> Any:
    from langfuse.langchain import CallbackHandler
    return CallbackHandler


def _has_langfuse_credentials() -> bool:
    return bool(
        os.getenv("LANGFUSE_PUBLIC_KEY", "").strip()
        and os.getenv("LANGFUSE_SECRET_KEY", "").strip()
    )


def _start_workflow_observation(
    *,
    name: str,
    input: dict[str, Any] | None = None,
    metadata: dict[str, Any] | None = None,
    parent_span_id: str | None = None,
) -> Any:
    if not _has_langfuse_credentials():
        return nullcontext(None)
    try:
        return _import_langfuse_get_client()().start_as_current_observation(
            name=name,
            as_type="span",
            input=input,
            metadata=metadata,
        )
    except Exception:
        logger.debug(
            "langfuse.workflow_start_observation_failed name=%s",
            name,
            exc_info=True,
        )
        return nullcontext(None)


def _update_workflow_observation(
    observation: Any | None,
    *,
    output: dict[str, Any] | None = None,
    level: str | None = None,
    status_message: str | None = None,
) -> None:
    if observation is None:
        return
    try:
        kwargs: dict[str, Any] = {}
        if output is not None:
            kwargs["output"] = output
        if level is not None:
            kwargs["level"] = level
        if status_message is not None:
            kwargs["status_message"] = status_message
        observation.update(**kwargs)
    except Exception:
        logger.debug("langfuse.workflow_update_failed", exc_info=True)


def _build_workflow_llm_callbacks(
    *,
    parent_span_id: str | None = None,
) -> tuple[Any, ...]:
    if not _has_langfuse_credentials():
        return ()
    if not langfuse_context_has_active_span():
        return ()

    try:
        client = _import_langfuse_get_client()()
        trace_id = str(client.get_current_trace_id() or "").strip()
        if not trace_id:
            return ()
        trace_context: dict[str, str] = {"trace_id": trace_id}
        if parent_span_id:
            trace_context["parent_span_id"] = parent_span_id
        return (_import_langfuse_callback_handler()(trace_context=trace_context),)
    except Exception:
        logger.debug("langfuse.workflow_callback_handler_failed", exc_info=True)
        return ()


def _summarize_mapping_keys(value: Any) -> list[str]:
    if not isinstance(value, dict):
        return []
    return sorted(str(key) for key in value.keys())


def _action_observation_input(
    *,
    params: dict[str, Any],
) -> dict[str, Any]:
    return {
        "input_keys": _summarize_mapping_keys(params),
    }


def _action_observation_output(
    *,
    result: ActionResult,
    next_target: str | None = None,
    paused: bool = False,
    pause_def: PauseDef | None = None,
) -> dict[str, Any]:
    output: dict[str, Any] = {
        "success": result.success,
        "output_keys": _summarize_mapping_keys(result.output),
        "metadata_keys": _summarize_mapping_keys(result.metadata),
    }
    if isinstance(result.output, dict):
        human_validation_text = result.output.get("texto_validacao_humana")
        if isinstance(human_validation_text, str) and human_validation_text.strip():
            output["texto_validacao_humana"] = human_validation_text.strip()
    if result.error:
        output["error"] = result.error
    if next_target is not None:
        output["next_target"] = next_target
    if paused and pause_def is not None:
        output["status"] = "WAITING_INPUT"
        output["resume_from"] = pause_def.resume_from
        output["expected_input_key"] = pause_def.expected_input.key
    return output


def _workflow_metadata(
    *,
    workflow_name: str,
    workflow_version: int,
    node_id: str | None = None,
    action_name: str | None = None,
    label: str | None = None,
    stage: str | None = None,
) -> dict[str, Any]:
    metadata: dict[str, Any] = {
        "workflow_name": workflow_name,
        "workflow_version": str(workflow_version),
    }
    if node_id:
        metadata["node_id"] = node_id
    if action_name:
        metadata["action"] = action_name
    if label:
        metadata["label"] = label
    if stage:
        metadata["stage"] = stage
    return metadata


class WorkflowState(TypedDict, total=False):
    input: dict[str, Any]
    vars: dict[str, Any]
    trace: list[dict[str, Any]]
    status: str
    current_node: str | None
    last_node: str | None
    pending_interrupt: dict[str, Any] | None
    final_data: Any
    final_error: str | None
    final_metadata: dict[str, Any]


def compile_workflow(
    definition: WorkflowDef,
    *,
    action_registry: ActionRegistry,
    runtime: WorkflowRuntimeContext,
    checkpointer: Any,
) -> Any:
    try:
        from langgraph.graph import END, START, StateGraph
        from langgraph.types import Command
    except ModuleNotFoundError as exc:  # pragma: no cover - depende do ambiente
        raise ModuleNotFoundError(
            "langgraph nao esta instalado. Adicione a dependencia para usar workflows."
        ) from exc

    nodes_by_id = {node.id: node for node in definition.nodes}
    edges_by_source: dict[str, list[EdgeDef]] = {}
    for edge in sorted(definition.edges, key=lambda item: item.priority):
        edges_by_source.setdefault(edge.source, []).append(edge)

    builder = StateGraph(WorkflowState)

    def finish_node(state: WorkflowState) -> dict[str, Any]:
        # Preserve the terminal state set by the last workflow step.
        # Returning an empty payload here can drop status/final_data in some
        # graph runtimes/checkpointer combinations, making the service read
        # a FAILED default at the end.
        return dict(state)

    builder.add_node(FINISH_NODE, finish_node)
    builder.add_edge(FINISH_NODE, END)

    for node in definition.nodes:
        builder.add_node(
            node.id,
            _make_action_node(
                definition=definition,
                node=node,
                outgoing_edges=edges_by_source.get(node.id, []),
                nodes_by_id=nodes_by_id,
                action_registry=action_registry,
                runtime=runtime,
                command_type=Command,
            ),
        )
        if node.pause is not None and node.pause.enabled:
            builder.add_node(
                _pause_node_name(node.id),
                _make_pause_node(
                    node=node,
                    command_type=Command,
                ),
            )

    builder.add_edge(START, definition.start)
    return builder.compile(checkpointer=checkpointer)


def build_initial_state(
    definition: WorkflowDef, input_payload: dict[str, Any]
) -> WorkflowState:
    return WorkflowState(
        input=deepcopy(input_payload),
        vars={},
        trace=[],
        status="RUNNING",
        current_node=definition.start,
        last_node=None,
        pending_interrupt=None,
        final_data=None,
        final_error=None,
        final_metadata={
            "workflow": definition.name,
            "version": definition.version,
        },
    )


def _make_action_node(
    *,
    definition: WorkflowDef,
    node: NodeDef,
    outgoing_edges: list[EdgeDef],
    nodes_by_id: dict[str, NodeDef],
    action_registry: ActionRegistry,
    runtime: WorkflowRuntimeContext,
    command_type: Any,
):
    def action_node(state: WorkflowState) -> Any:
        context = _build_context(state)
        params = render_template(node.input, context)

        try:
            handler = action_registry.get(node.action)
        except ValueError as exc:
            raise WorkflowConfigurationError(str(exc)) from exc

        action_label = _ACTION_LABELS.get(
            node.action, node.action,
        )
        logger.info(
            "workflow.iniciou | %s",
            action_label,
        )
        step_metadata = _workflow_metadata(
            workflow_name=definition.name,
            workflow_version=definition.version,
            node_id=node.id,
            action_name=node.action,
            label=action_label,
            stage="step",
        )
        with _start_workflow_observation(
            name=f"workflow.step.{node.id}",
            input=_action_observation_input(params=params),
            metadata=step_metadata,
        ) as step_observation:
            emit_workflow_step({
                "stage": "step_started",
                "step": node.id,
                "label": action_label,
            })
            step_parent_span_id = (
                str(getattr(step_observation, "id", "")).strip() or None
            )
            step_runtime = replace(
                runtime,
                llm_callbacks=_build_workflow_llm_callbacks(
                    parent_span_id=step_parent_span_id,
                ),
                llm_metadata=step_metadata,
            )
            try:
                action_result = handler(
                    state, params, step_runtime,
                )
            except Exception as exc:
                _update_workflow_observation(
                    step_observation,
                    output={"error": str(exc)},
                    level="ERROR",
                    status_message=str(exc),
                )
                raise
            updated_state = _apply_action_result(
                state, node.id, node.action,
                action_result,
            )
            logger.info(
                "workflow.finalizou | %s | success=%s",
                action_label, action_result.success,
            )
            emit_workflow_step({
                "stage": "step_completed",
                "step": node.id,
                "label": action_label,
                "success": action_result.success,
            })

            if not action_result.success:
                logger.error(
                    "workflow.step.failed workflow=%s version=%s node=%s "
                    "action=%s error=%s metadata=%s",
                    definition.name,
                    definition.version,
                    node.id,
                    node.action,
                    action_result.error,
                    action_result.metadata,
                )
                updated_state["status"] = "FAILED"
                updated_state["current_node"] = None
                updated_state["final_data"] = None
                updated_state["final_error"] = action_result.error
                updated_state["pending_interrupt"] = None
                updated_state["final_metadata"] = {
                    "workflow": definition.name,
                    "version": definition.version,
                    "failed_node": node.id,
                    **action_result.metadata,
                }
                _update_workflow_observation(
                    step_observation,
                    output=_action_observation_output(
                        result=action_result,
                        next_target=FINISH_NODE,
                    ),
                    level="ERROR",
                    status_message=action_result.error or "workflow step failed",
                )
                return command_type(update=updated_state, goto=FINISH_NODE)

            if node.pause is not None and node.pause.enabled:
                pause_state = _build_pause_state(
                    definition=definition,
                    state=updated_state,
                    node=node,
                    result=action_result,
                    parent_span_id=step_parent_span_id,
                )
                if pause_state is not None:
                    _update_workflow_observation(
                        step_observation,
                        output=_action_observation_output(
                            result=action_result,
                            next_target=_pause_node_name(node.id),
                            paused=True,
                            pause_def=node.pause,
                        ),
                    )
                    return command_type(
                        update=pause_state,
                        goto=_pause_node_name(node.id),
                    )

            target = _resolve_next_target(
                outgoing_edges,
                updated_state,
                workflow_name=definition.name,
                workflow_version=definition.version,
                parent_span_id=step_parent_span_id,
            )
            if target is None or target == "END":
                logger.info("workflow.end node=%s", node.id)
                updated_state["status"] = "COMPLETED"
                updated_state["current_node"] = None
                updated_state["final_data"] = action_result.output
                updated_state["final_error"] = None
                updated_state["pending_interrupt"] = None
                updated_state["final_metadata"] = {
                    "workflow": definition.name,
                    "version": definition.version,
                    "last_node": node.id,
                    **action_result.metadata,
                }
                _update_workflow_observation(
                    step_observation,
                    output=_action_observation_output(
                        result=action_result,
                        next_target="END",
                    ),
                )
                return command_type(update=updated_state, goto=FINISH_NODE)

            if target not in nodes_by_id:
                raise WorkflowConfigurationError(
                    f"Destino {target!r} nao encontrado para o node {node.id!r}"
                )

            updated_state["status"] = "RUNNING"
            updated_state["current_node"] = target
            updated_state["pending_interrupt"] = None
            _update_workflow_observation(
                step_observation,
                output=_action_observation_output(
                    result=action_result,
                    next_target=target,
                ),
            )
            return command_type(update=updated_state, goto=target)

    return action_node


def _make_pause_node(*, node: NodeDef, command_type: Any):
    pause_def = node.pause
    assert pause_def is not None

    def pause_node(state: WorkflowState) -> Any:
        try:
            from langgraph.errors import GraphInterrupt
            from langgraph.types import interrupt
        except ModuleNotFoundError as exc:  # pragma: no cover - depende do ambiente
            raise ModuleNotFoundError(
                "langgraph nao esta instalado. "
                "Adicione a dependencia para usar workflows."
            ) from exc

        pending = state.get("pending_interrupt")
        if not isinstance(pending, dict):
            raise WorkflowConfigurationError(
                f"Node de pausa {node.id!r} sem pending_interrupt no estado"
            )

        graph_interrupt: GraphInterrupt | None = None
        with _start_workflow_observation(
            name="workflow.resume",
            input={
                "node_id": node.id,
                "expected_input_key": pause_def.expected_input.key,
                "resume_from": pause_def.resume_from,
            },
            metadata={
                "node_id": node.id,
                "resume_from": pause_def.resume_from,
                "stage": "resume",
            },
        ) as resume_observation:
            input_date: dict[str, Any] | None = None
            try:
                resume_payload = interrupt(pending["payload"])
                input_date = dict(state.get("input", {}))
                resumed_input = _resume_input_dict(
                    resume_payload,
                    pause_def.expected_input.key,
                )
                input_date.update(resumed_input)
                _normalize_and_validate_input(input_date, pause_def)
            except GraphInterrupt as exc:
                graph_interrupt = exc
                _update_workflow_observation(
                    resume_observation,
                    output={
                        "status": "WAITING_INPUT",
                        "resume_from": pause_def.resume_from,
                        "expected_input_key": pause_def.expected_input.key,
                    },
                )
            except Exception as exc:
                _update_workflow_observation(
                    resume_observation,
                    output={"error": str(exc)},
                    level="ERROR",
                    status_message=str(exc),
                )
                raise

            if graph_interrupt is None:
                if input_date is None:
                    raise WorkflowConfigurationError(
                        f"Node de pausa {node.id!r} nao recebeu input de retomada"
                    )
                updated_state = _clone_state(state)
                updated_state["input"] = input_date
                updated_state["status"] = "RUNNING"
                updated_state["pending_interrupt"] = None
                updated_state["current_node"] = pause_def.resume_from
                updated_state["trace"].append(
                    {
                        "node_id": node.id,
                        "action": "resume_input",
                        "success": True,
                        "error": None,
                    }
                )

                if pause_def.resume_from == "END":
                    updated_state["status"] = "COMPLETED"
                    updated_state["current_node"] = None
                    updated_state["final_data"] = updated_state.get("vars", {}).get(node.id)
                    updated_state["final_error"] = None
                    _update_workflow_observation(
                        resume_observation,
                        output={
                            "resume_from": "END",
                            "normalized_input_keys": _summarize_mapping_keys(input_date),
                        },
                    )
                    return command_type(update=updated_state, goto=FINISH_NODE)

                _update_workflow_observation(
                    resume_observation,
                    output={
                        "resume_from": pause_def.resume_from,
                        "normalized_input_keys": _summarize_mapping_keys(input_date),
                    },
                )
                return command_type(update=updated_state, goto=pause_def.resume_from)

        if graph_interrupt is not None:
            raise graph_interrupt
        raise WorkflowConfigurationError(
            f"Node de pausa {node.id!r} finalizou sem estado valido de retomada"
        )

    return pause_node


def _apply_action_result(
    state: WorkflowState,
    node_id: str,
    action_name: str,
    result: ActionResult,
) -> WorkflowState:
    updated_state = _clone_state(state)
    vars_state = dict(updated_state.get("vars", {}))
    vars_state[node_id] = deepcopy(result.output)
    updated_state["vars"] = vars_state

    trace = list(updated_state.get("trace", []))
    trace.append(
        {
            "node_id": node_id,
            "action": action_name,
            "success": result.success,
            "error": result.error,
        }
    )
    updated_state["trace"] = trace
    updated_state["last_node"] = node_id
    updated_state["current_node"] = node_id
    return updated_state


def _build_pause_state(
    *,
    definition: WorkflowDef,
    state: WorkflowState,
    node: NodeDef,
    result: ActionResult,
    parent_span_id: str | None = None,
) -> WorkflowState | None:
    pause_def = node.pause
    if pause_def is None:
        return None

    pause_context = _build_context(state, output=result.output)
    if not evaluate_condition(pause_def.when, pause_context):
        return None

    payload = resolve_value(pause_def.return_from, pause_context)
    if payload is None:
        raise WorkflowConfigurationError(
            "pause.return_from="
            f"{pause_def.return_from!r} nao resolveu valor "
            f"no node {node.id!r}"
        )

    updated_state = _clone_state(state)
    updated_state["status"] = "WAITING_INPUT"
    updated_state["current_node"] = node.id
    updated_state["pending_interrupt"] = {
        "node_id": node.id,
        "payload": payload,
        "resume_from": pause_def.resume_from,
        "expected_input_key": pause_def.expected_input.key,
        "allowed_values": list(pause_def.expected_input.allowed_values),
        "normalize": pause_def.expected_input.normalize,
    }
    updated_state["final_metadata"] = {
        "workflow": definition.name,
        "version": definition.version,
        "paused_at": node.id,
        "resume_from": pause_def.resume_from,
        "expected_input_key": pause_def.expected_input.key,
        "allowed_values": list(pause_def.expected_input.allowed_values),
        "normalize": pause_def.expected_input.normalize,
        **result.metadata,
    }
    with _start_workflow_observation(
        name="workflow.pause",
        input={
            "node_id": node.id,
            "resume_from": pause_def.resume_from,
            "expected_input_key": pause_def.expected_input.key,
        },
        metadata=_workflow_metadata(
            workflow_name=definition.name,
            workflow_version=definition.version,
            node_id=node.id,
            action_name=node.action,
            stage="pause",
        ),
        parent_span_id=parent_span_id,
    ) as pause_observation:
        _update_workflow_observation(
            pause_observation,
            output={
                "resume_from": pause_def.resume_from,
                "expected_input_key": pause_def.expected_input.key,
                "allowed_values": list(pause_def.expected_input.allowed_values),
            },
        )
    return updated_state


def _resolve_next_target(
    outgoing_edges: list[EdgeDef],
    state: WorkflowState,
    *,
    workflow_name: str = "",
    workflow_version: int = 0,
    parent_span_id: str | None = None,
) -> str | None:
    context = _build_context(state)
    current = state.get("current_node", "?")

    # Extrair info do perfil de fatura para logs descritivos
    vars_state = context.get("vars", {})
    perfil = vars_state.get("consultar_perfil_fatura", {})
    forma_pagamento = ""
    if isinstance(perfil, dict):
        forma_pagamento = str(
            perfil.get("forma_pagamento", "")
        ).strip()

    for edge in outgoing_edges:
        matched = evaluate_condition(edge.when, context)
        if matched:
            decision = ""
            # Decisões de fatura
            if (
                current == "consultar_perfil_fatura"
                and forma_pagamento
            ):
                if edge.target == "registrar_protocolo":
                    decision = (
                        f"Fatura tipo {forma_pagamento}"
                        " → crédito na próxima fatura"
                    )
                else:
                    decision = (
                        f"Fatura tipo {forma_pagamento}"
                        " → verificando se está aberta"
                    )
            elif current == "check_invoice_status":
                if edge.target == "enviar_sms":
                    decision = (
                        "Fatura aberta"
                        " → enviando SMS com boleto"
                    )
                else:
                    decision = (
                        "Fatura fechada"
                        " → crédito na próxima fatura"
                    )

            if decision:
                logger.info(
                    "workflow.decisao | %s", decision,
                )
                emit_workflow_step({
                    "stage": "decision",
                    "from": current,
                    "to": edge.target,
                    "label": decision,
                })
            else:
                logger.info(
                    "workflow.route | %s → %s",
                    current, edge.target,
                )
            with _start_workflow_observation(
                name="workflow.decision",
                input={
                    "from": current,
                    "to": edge.target,
                },
                metadata=_workflow_metadata(
                    workflow_name=workflow_name or "workflow",
                    workflow_version=workflow_version or 0,
                    node_id=str(current),
                    stage="decision",
                    label=decision or f"{current} -> {edge.target}",
                ),
                parent_span_id=parent_span_id,
            ) as decision_observation:
                _update_workflow_observation(
                    decision_observation,
                    output={
                        "matched": True,
                        "from": current,
                        "to": edge.target,
                        "label": decision or "",
                    },
                )
            return edge.target
    return None


def _normalize_and_validate_input(
    input_date: dict[str, Any], pause_def: PauseDef
) -> None:
    expected = pause_def.expected_input
    if expected.key not in input_date:
        raise WorkflowInputError(
            f"Input obrigatorio ausente para continuar fluxo: {expected.key}"
        )

    value = input_date[expected.key]
    normalized = _normalize_input_value(value, expected.normalize)
    input_date[expected.key] = normalized

    if expected.allowed_values:
        allowed = {
            str(_normalize_input_value(item, expected.normalize))
            for item in expected.allowed_values
        }
        if str(normalized) not in allowed:
            raise WorkflowInputError(
                f"Valor invalido para {expected.key}: {normalized!r}. "
                f"Permitidos: {sorted(allowed)}"
            )


def _normalize_input_value(value: Any, normalize_mode: str | None) -> Any:
    if not isinstance(value, str) or normalize_mode is None:
        return value
    if normalize_mode == "strip":
        return value.strip()
    if normalize_mode == "upper":
        return value.upper()
    if normalize_mode == "lower":
        return value.lower()
    if normalize_mode == "upper_strip":
        return value.strip().upper()
    if normalize_mode == "lower_strip":
        return value.strip().lower()
    return value


def _resume_input_dict(resume_payload: Any, expected_input_key: str) -> dict[str, Any]:
    if isinstance(resume_payload, dict):
        return dict(resume_payload)
    return {expected_input_key: resume_payload}


def _build_context(
    state: WorkflowState,
    *,
    output: dict[str, Any] | None = None,
) -> dict[str, Any]:
    return {
        "input": deepcopy(state.get("input", {})),
        "vars": deepcopy(state.get("vars", {})),
        "trace": deepcopy(state.get("trace", [])),
        "status": state.get("status"),
        "current_node": state.get("current_node"),
        "last_node": state.get("last_node"),
        "output": deepcopy(output or {}),
    }


def _clone_state(state: WorkflowState) -> WorkflowState:
    return WorkflowState(
        input=deepcopy(state.get("input", {})),
        vars=deepcopy(state.get("vars", {})),
        trace=deepcopy(state.get("trace", [])),
        status=str(state.get("status", "RUNNING")),
        current_node=state.get("current_node"),
        last_node=state.get("last_node"),
        pending_interrupt=deepcopy(state.get("pending_interrupt")),
        final_data=deepcopy(state.get("final_data")),
        final_error=state.get("final_error"),
        final_metadata=deepcopy(state.get("final_metadata", {})),
    )


def _pause_node_name(node_id: str) -> str:
    return f"__pause__{node_id}"