from __future__ import annotations

import secrets
import uuid
from dataclasses import dataclass, field
from datetime import datetime, timezone
from enum import Enum
from typing import Any, Awaitable, Callable, Dict, List, Optional, Union

import logging

logger = logging.getLogger(__name__)


class NodeStatus(Enum):
    PENDING = "pending"
    IN_PROGRESS = "in_progress"
    SUCCESS = "success"
    FAILED = "failed"
    RETRY = "retry"


@dataclass
class TalkbackTraceStep:
    name: str
    status: NodeStatus = NodeStatus.PENDING
    duration_ms: Optional[int] = None
    error: Optional[str] = None


@dataclass
class TalkbackTrace:
    trace_id: str
    steps: Dict[str, TalkbackTraceStep] = field(default_factory=dict)

    def start(self, name: str) -> None:
        self.steps[name] = TalkbackTraceStep(name=name, status=NodeStatus.IN_PROGRESS)

    def finish(self, name: str, status: NodeStatus, duration_ms: int, error: Optional[str] = None) -> None:
        step = self.steps.get(name) or TalkbackTraceStep(name=name)
        step.status = status
        step.duration_ms = duration_ms
        step.error = error
        self.steps[name] = step


@dataclass
class TalkbackState:
    # Inputs
    request: Any
    current_user: Any
    current_user_id: str

    # Derived / working
    question_text: str = ""
    resolved_agent_ids: List[str] = field(default_factory=list)
    dispatch_agent_ids: List[str] = field(default_factory=list)
    datasource: Optional[str] = None

    service: Any = None
    session_key: Optional[str] = None
    session_state: Optional[Dict[str, Any]] = None
    history_turns: List[Dict[str, Any]] = field(default_factory=list)
    session_context: Dict[str, Any] = field(default_factory=dict)

    context_cleared: bool = False
    context_clear_reason: str = ""

    routing_question_text: str = ""
    intent_routing: Optional[Dict[str, Any]] = None
    # Set in LoadSession when clarification_context stores clarification_agent_ids (narrow routing)
    routing_agent_scope_ids: Optional[List[str]] = None
    # Set in LoadSession when client opts in (search_web) after llm_prior_df consent
    pending_prior_df_followup: Optional[str] = None
    # Previous agent_scope from session state — used by DispatchDecisionNode to auto-select
    # the common agent when new routing overlaps with the prior turn's scope.
    previous_agent_scope: Optional[List[str]] = None
    # Snapshot of the prior turn's dispatch scope captured before context-clear wipes it.
    # DispatchDecisionNode falls back to this when state.context_cleared so identical or
    # subset-of-prior follow-ups can reuse the previous agents without re-asking the user.
    pre_clear_dispatch_scope: Optional[List[str]] = None
    file_cmd_parse: Optional[Dict[str, Any]] = None
    # User answered agent/source disambiguation this request (LoadSession appended synthetic history row).
    agent_disambiguation_follow_up: bool = False
    # Short reply before merge (e.g. "Order Management, AI agent") — for routing prompt only.
    disambiguation_user_choice: Optional[str] = None
    # When set, ``append_turn_to_talkback_session`` prefers this as the stored ``question`` (original ask).
    talkback_history_question_text: Optional[str] = None
    # Verbatim question from the client at request start (before ordinal resolve / non-file strip).
    talkback_raw_question_text: str = ""
    # Client thread id: request ``session_key`` if non-empty, else server-minted ``uid::utc::random``.
    talkback_effective_session_key: str = ""
    # One-shot skip after user answers the broad-scope clarification prompt (session-driven).
    skip_broad_scope_check_once: bool = False
    # One-shot skip after user chose "keep current context" on out-of-context clear confirmation.
    skip_context_relevance_once: bool = False

    result: Optional[Dict[str, Any]] = None
    response: Optional[Dict[str, Any]] = None

    # Set by UserQuestionLog cache hit: handler already ran; pipeline skips routing through Execute.
    _uql_cache_hit_short_circuit: bool = False

    trace: TalkbackTrace = field(default_factory=lambda: TalkbackTrace(trace_id=str(uuid.uuid4())))


# Dependency contracts (kept generic to avoid circular imports)
ResolveAuthorizedAgentIds = Callable[[Optional[Union[str, List[str]]], str], List[str]]
GetAskmeService = Callable[[Any], Any]
RunInThreadpool = Callable[..., Awaitable[Any]]
ExecuteTalkback = Callable[[Any, Any], Awaitable[Dict[str, Any]]]
