from __future__ import annotations

import json
from datetime import datetime, timedelta
from typing import TYPE_CHECKING, Any, Iterable, Mapping

import rich
from rich import box as rich_box
from rich.console import Console
from rich.panel import Panel
from rich.table import Table

if TYPE_CHECKING:
    from relationalai.config.config import Config
    from relationalai.config.connections import ConnectionConfig

from relationalai.services.reasoners.models import ReasonerType


def format_state_with_color(state: str) -> str:
    """Format state with colors (mirrors upstream CLI `format_state_with_color`)."""
    if not state:
        return ""
    state_upper = state.upper()
    # Success / terminal-ok states (jobs + operations).
    if state_upper in ("COMPLETED", "SUCCEEDED"):
        return f"[green]{state_upper}[/green]"
    if state_upper == "READY":
        return f"[green]{state_upper}[/green]"
    if state_upper == "SUSPENDED":
        return f"[yellow]{state_upper}[/yellow]"
    if state_upper in ("CANCELED", "CANCELLED"):
        return f"[yellow]{state_upper}[/yellow]"
    if state_upper in ("PENDING", "SYNCING", "PROCESSING", "RUNNING", "CANCELLING"):
        return f"[bold yellow]{state_upper}[/bold yellow]"
    if state_upper in ("ABORTED", "QUARANTINED", "GONE", "FAILED", "ERROR", "DELETED"):
        return f"[red]{state_upper}[/red]"
    return state_upper


def format_value(value: Any) -> str:
    """Format values for table display (mirrors upstream CLI `format_value`)."""
    if value is None:
        return "N/A"
    if isinstance(value, (int, float, bool)):
        return f"{value}"
    if isinstance(value, list):
        return ", ".join(map(str, value))
    if isinstance(value, datetime):
        return value.strftime("%Y-%m-%d %H:%M:%S")
    if isinstance(value, timedelta):
        # Keep it simple here; reasoners tables typically don't show timedeltas.
        return str(value)
    return str(value)


def _format_duration_ms(duration_ms: int | float | None) -> str:
    """Format a duration in milliseconds into a compact human form."""
    if duration_ms is None:
        return "N/A"
    try:
        ms = float(duration_ms)
    except Exception:
        return str(duration_ms)
    if ms < 1000:
        return f"{int(ms)}ms" if ms.is_integer() else f"{ms}ms"
    s = ms / 1000.0
    # Keep one decimal only when needed (e.g. 1.5s), otherwise show integer seconds.
    if abs(s - round(s)) < 1e-9:
        return f"{int(round(s))}s"
    out = f"{s:.1f}".rstrip("0").rstrip(".")
    return f"{out}s"


_KNOWN_REASONER_TYPES = {ReasonerType.LOGIC, ReasonerType.PRESCRIPTIVE, ReasonerType.PREDICTIVE}


def _get(obj: Any, key: str, default: Any = "") -> Any:
    if isinstance(obj, Mapping):
        return obj.get(key, default)
    return getattr(obj, key, default)


def _format_reasoner_type(reasoner_type: str) -> str:
    t = (reasoner_type or "").upper()
    from relationalai.tools.cli.job_types import JOB_TYPE_LABELS

    if t in JOB_TYPE_LABELS:
        return JOB_TYPE_LABELS[t]
    if t in _KNOWN_REASONER_TYPES:
        return ReasonerType.get_label(t)
    return reasoner_type or ""


def show_reasoners(reasoners: Iterable[Any]) -> None:
    """Show reasoners in a table matching upstream reasoners table layout."""
    items = list(reasoners)
    if not len(items):
        return

    table = Table(show_header=True, border_style="dim", header_style="bold", box=rich_box.SIMPLE_HEAD)
    table.add_column("#")
    table.add_column("Name")
    table.add_column("Type")
    table.add_column("Size")
    table.add_column("State")
    table.add_column("Created By")
    table.add_column("Created On")

    for index, r in enumerate(items):
        r_type = _get(r, "type", "") or ""
        type_display = _format_reasoner_type(str(r_type))
        created_on = format_value(_get(r, "created_on", None))
        state_display = format_state_with_color(str(_get(r, "state", "") or ""))
        table.add_row(
            f"{index+1}",
            str(_get(r, "name", "")),
            type_display,
            str(_get(r, "size", "") or ""),
            state_display,
            str(_get(r, "created_by", "") or ""),
            created_on,
        )

    rich.print(table)


def show_reasoner(reasoner: Any) -> None:
    """Convenience wrapper for printing a single reasoner."""
    show_reasoners([reasoner])


def show_jobs(jobs: Iterable[Any], *, include_details: bool = False) -> None:
    """Show jobs in a stable, readable table.

    When `include_details=True`, shows extra columns populated from `job.details`
    (if present) such as database/read_only/timeout_ms.
    """
    items = list(jobs)
    if not len(items):
        return

    table = Table(show_header=True, border_style="dim", header_style="bold", box=rich_box.SIMPLE_HEAD)
    table.add_column("#")
    table.add_column("ID", no_wrap=True)
    table.add_column("Type")
    table.add_column("State")
    table.add_column("Reasoner")
    table.add_column("Created By")
    table.add_column("Created On")
    table.add_column("Finished At")
    table.add_column("Duration")
    table.add_column("Abort Reason")
    if include_details:
        table.add_column("Database")
        table.add_column("Read Only")
        table.add_column("Timeout (ms)")

    for index, j in enumerate(items):
        t_display = _format_reasoner_type(str(_get(j, "reasoner_type", "") or ""))
        state_display = format_state_with_color(str(_get(j, "state", "") or ""))
        details = _get(j, "details", None)
        details_d: dict[str, Any] = dict(details) if isinstance(details, Mapping) else {}
        dur_any = _get(j, "duration", None)
        duration_display = (
            _format_duration_ms(dur_any)
            if (dur_any is None or isinstance(dur_any, (int, float)))
            else format_value(dur_any)
        )
        table.add_row(
            f"{index+1}",
            str(_get(j, "id", "")),
            t_display,
            state_display,
            format_value(_get(j, "name", None)),
            format_value(_get(j, "created_by", None)),
            format_value(_get(j, "created_on", None)),
            format_value(_get(j, "finished_at", None)),
            duration_display,
            format_value(_get(j, "abort_reason", None)),
            *(
                [
                    format_value(details_d.get("database_name")),
                    format_value(details_d.get("read_only")),
                    format_value(details_d.get("timeout_ms")),
                ]
                if include_details
                else []
            ),
        )

    rich.print(table)


def show_job(job: Any) -> None:
    """Convenience wrapper for printing a single job."""
    show_jobs([job])


def show_connect_panel(config: Config, console: Console) -> ConnectionConfig:
    """Print the rai connect info panel and return the resolved default connection."""
    from relationalai.tools.cli.error_handlers import format_connection_info

    connection = config.get_default_connection()

    info_table = Table(show_header=False, box=None, padding=(0, 2, 0, 0), expand=True)
    info_table.add_column(style="dim cyan", no_wrap=True)
    info_table.add_column(style="white")

    source_field = type(config).model_fields.get("source")
    config_source = source_field.default if source_field is not None else None
    if config_source:
        info_table.add_row("Config", config_source)
    if config.active_profile:
        info_table.add_row("Profile", config.active_profile)

    if config_source or config.active_profile:
        info_table.add_row("", "")  # spacer

    for label, value in format_connection_info(connection):
        info_table.add_row(label, value)

    console.print(Panel(info_table, title="[bold cyan]rai connect[/bold cyan]", border_style="cyan", padding=(1, 2)))
    return connection


def print_job_expansions(job: Any) -> None:
    """Print expanded job sections when present (mirrors scripts/jobs.py UX)."""
    if getattr(job, "details", None) is not None:
        rich.print("\n[bold]DETAILS[/bold]")
        rich.print(json.dumps(getattr(job, "details"), indent=2, default=str))
    if getattr(job, "artifacts", None) is not None:
        rich.print("\n[bold]ARTIFACTS[/bold]")
        rich.print(json.dumps(getattr(job, "artifacts"), indent=2, default=str))
    if getattr(job, "problems", None) is not None:
        rich.print("\n[bold]PROBLEMS[/bold]")
        rich.print(json.dumps(getattr(job, "problems"), indent=2, default=str))
    if getattr(job, "events", None) is not None:
        ev = getattr(job, "events")
        rich.print("\n[bold]EVENTS[/bold]")
        rich.print(
            json.dumps(
                {"events": getattr(ev, "events", []), "continuation_token": getattr(ev, "continuation_token", None)},
                indent=2,
                default=str,
            )
        )
    if getattr(job, "raw", None) is not None:
        rich.print("\n[bold]RAW[/bold]")
        rich.print(json.dumps(getattr(job, "raw"), indent=2, default=str))

