from __future__ import annotations

import time
from typing import Dict, Optional, TYPE_CHECKING

from relationalai import debugging
from relationalai.clients.util import poll_with_specified_overhead
from relationalai.clients.config import Config
from relationalai.tools.cli_controls import create_progress
from relationalai.tools.txn_progress_tree import build_progress_tree, compute_progress_summary
from relationalai.tools.txn_progress_display import format_progress_summary
from relationalai.util.format import format_duration
from relationalai.tools.constants import USE_DIRECT_ACCESS
from relationalai.tools.txn_progress_display import TxnProgressUpdater

if TYPE_CHECKING:
    from relationalai.clients.resources.snowflake import Resources
    from relationalai.clients.resources.snowflake.snowflake import TxnStatusResponse

# Polling behavior constants
POLL_OVERHEAD_RATE = 0.1  # Overhead rate for exponential backoff
MAX_POLL_DELAY_SECONDS = 20
MAX_POLL_DELAY_FOR_DIRECT_ACCESS_SECONDS = 10

# Text color constants
GREEN_COLOR = '\033[92m'
GRAY_COLOR = '\033[90m'
ENDC = '\033[0m'

PRINT_TXN_PROGRESS_FLAG = "print_txn_progress"
PRINT_INTERNAL_TXN_PROGRESS_FLAG = "print_txn_progress_internal"

def should_print_txn_progress(config: Config) -> bool:
    return bool(config.get(PRINT_TXN_PROGRESS_FLAG, True))

def should_print_internal_txn_progress(config) -> bool:
    return bool(config.get(PRINT_INTERNAL_TXN_PROGRESS_FLAG, False))

class ExecTxnPoller:
    """
    Encapsulates the polling logic for exec_async transaction completion.
    """

    def __init__(
        self,
        config: Config,
        engine: str,
        model: str,
        resource: Optional["Resources"] = None,
        txn_id: Optional[str] = None,
        headers: Optional[Dict] = None,
        txn_start_time: Optional[float] = None,
        source_info: Optional[Dict] = None,
        suppress_creation_errors: bool = False,
    ):
        self.print_txn_progress = should_print_txn_progress(config)
        self.print_internal_txn_progress = should_print_internal_txn_progress(config)
        self.res = resource
        self.txn_id = txn_id
        self.headers = headers or {}
        self.engine = engine
        self.model = model
        self.source_info = source_info or {}
        self.txn_start_time = txn_start_time or time.time()
        self.last_status: Optional[TxnStatusResponse] = None
        self.max_poll_delay = MAX_POLL_DELAY_SECONDS
        if config.get("use_direct_access", USE_DIRECT_ACCESS):
            self.max_poll_delay = MAX_POLL_DELAY_FOR_DIRECT_ACCESS_SECONDS
        self.txn_progress_updater: Optional[TxnProgressUpdater] = None
        self.suppress_creation_errors = suppress_creation_errors
        self.evaluation_progress = {
            'completed_relations': 0,
        }

    def __enter__(self) -> ExecTxnPoller:
        if not self.print_txn_progress:
            return self
        print(self._format_source_label())

        self.progress = create_progress(
            description=lambda: self.description_with_timing(),
            success_message="",  # We'll handle this ourselves
            leading_newline=False,
            trailing_newline=False,
            show_duration_summary=False,
        )
        self.progress.__enter__()
        # Initialize the progress updater to manage execution tree display
        self.txn_progress_updater = TxnProgressUpdater(
            self.progress,
            enabled=self.print_internal_txn_progress
        )
        return self

    def __exit__(self, exc_type, exc_value, traceback) -> None:
        if not self.print_txn_progress:
            return

        if exc_type is not None and self.txn_id is None and self.suppress_creation_errors:
            # Transaction creation failed before we got a txn_id, and the
            # caller indicated it will handle retries (e.g., _exec_with_gi_retry
            # retrying after "database not found" on first model access).
            # Silently dismiss the progress display — the retry path will
            # create a new progress on the next attempt.
            live = getattr(self.progress, 'live', None)
            if live is not None:
                live.transient = True
                live.stop()
                self.progress.live = None  # type: ignore[union-attr]
            # Prevent failure message output in both terminal and notebook
            # progress displays (TaskProgress checks self.live which is now
            # None; NotebookTaskProgress checks self.failure_message).
            self.progress.failure_message = None  # type: ignore[assignment]
            self.progress.__exit__(exc_type, exc_value, traceback)
            return

        # Update to success message with duration
        total_duration = time.time() - self.txn_start_time
        txn_id = self.txn_id
        summary = query_progress_message(self.evaluation_progress, True, total_duration, txn_id)
        if self.print_internal_txn_progress and self.last_status and self.last_status.progress:
            task_tree = build_progress_tree(self.last_status.progress["tasks"])
            summary += "\n" + format_progress_summary(compute_progress_summary(task_tree), total_duration)
        self.progress.update_main_status(summary)
        self.progress.__exit__(exc_type, exc_value, traceback)
        return

    def _get_last_progress(self) -> Optional[Dict]:
        """Get last transaction progress message if enabled and available."""
        if self.print_txn_progress and self.last_status:
            return self.last_status.progress
        return None

    def poll(self) -> bool:
        """
        Poll for transaction completion with interactive progress display.

        Returns:
            True if transaction completed successfully, False otherwise
        """
        if not self.txn_id:
            raise ValueError("Transaction ID must be provided for polling.")
        else:
            txn_id = self.txn_id

        if self.print_txn_progress:
            # Update the main status to include the new txn_id
            self.progress.update_main_status_fn(
                lambda: self.description_with_timing(txn_id),
            )

        # Don't show duration summary - we handle our own completion message
        def check_status() -> bool:
            """Check if transaction is complete."""
            if self.res is None:
                raise ValueError("Resource must be provided for polling.")
            self.last_status = self.res._check_exec_async_status(txn_id, headers=self.headers)
            # Update the execution tree display if enabled
            if self.txn_progress_updater:
                progress_data = self.last_status.progress if self.last_status else None
                self.txn_progress_updater.update(progress_data)
            self.update_evaluation_progress()
            return self.last_status.finished

        with debugging.span("wait", txn_id=txn_id):
            poll_with_specified_overhead(check_status, overhead_rate=POLL_OVERHEAD_RATE, max_delay=self.max_poll_delay)

        return True

    def _format_source_label(self) -> str:
        """Format source location as 'Query @ file:line' or 'Query' if no source info."""
        file = self.source_info.get("file")
        line = self.source_info.get("line")
        if file:
            import os
            basename = os.path.basename(file)
            return f"Query @ {basename}:{line}"
        else:
            return "Querying..."

    def description_with_timing(self, txn_id: str | None = None) -> str:
        elapsed = time.time() - self.txn_start_time
        return query_progress_message(self.evaluation_progress, False, elapsed, txn_id)

    def update_evaluation_progress(self):
        """Update self.evaluation_progress from the txn_progress returned by check exec status."""
        txn_progress = self._get_last_progress()
        if txn_progress is None:
            return (0,0)
        completed_relations = 0
        in_progress_relations = 0
        for task_key, task in txn_progress.get('tasks', dict()).items():
            # print(task_key, task)
            if 'end_time' in task:
                completed_relations += 1
            else:
                in_progress_relations += 1
        self.evaluation_progress['completed_relations'] = completed_relations
        self.evaluation_progress['in_progress_relations'] = in_progress_relations

def relations_progress_body(in_progress: int, completed: int) -> str:
    return (
        # TODO(PR): Decide on verbiage. Relations? Views?
        f"   Relations Materialized: {completed} completed" +
        (f" ({in_progress} in progress)" if in_progress > 0 else "")
    )

def query_progress_message(evaluation_progress: dict, completed: bool, duration: float, id: str | None = None) -> str:
    relations_completed = evaluation_progress.get('completed_relations', 0)
    relations_in_progress = evaluation_progress.get('in_progress_relations', 0)
    if completed:
        # Assume all in-progress are completed. (This can happen if the engine doesn't send
        # its final observabilty updates before we mark the txn completed.)
        relations_completed += relations_in_progress
        relations_in_progress = 0
        # # There will always be at least the select relation, but it usually finishes before
        # # the last observability update, so we artifically add that.
        # relations_completed = max(1, relations_completed)

    # Don't print sub-second decimals, because it updates too fast and is distracting.
    duration_str = format_duration(duration)
    if id is None:
        # Print with whitespace to align with the end of the transaction ID
        return f"Submitting job...                            {duration_str}"
    else:
        out = f"ID: {GRAY_COLOR}{id}{ENDC}  {duration_str}"

    if relations_completed + relations_in_progress > 0:
        out += f"\n{ENDC}  |  Evaluating relations: {relations_completed} completed"
        if relations_in_progress > 0:
            out += f" ({relations_in_progress} in progress)"
        out += ENDC

    return out
