import type { TemplateResult } from "lit";
import { css, html, LitElement, nothing } from "lit";
import { customElement, property } from "lit/decorators";
import { formatNumber } from "../../../../common/number/format_number";
import type { LocalizeKeys } from "../../../../common/translations/localize";
import "../../../../components/ha-alert";
import "../../../../components/ha-button";
import "../../../../components/ha-card";
import "../../../../components/ha-expansion-panel";
import "../../../../components/ha-spinner";
import "../../../../components/ha-yaml-editor";
import type { PipelineRun } from "../../../../data/assist_pipeline";
import type {
  ChatLog,
  ChatLogAssistantContent,
  ChatLogContent,
  ChatLogUserContent,
} from "../../../../data/chat_log";
import { showAlertDialog } from "../../../../dialogs/generic/show-dialog-box";
import type { HomeAssistant } from "../../../../types";

const RUN_DATA = ["pipeline", "language"];
const WAKE_WORD_DATA = ["engine"];

const STT_DATA = ["engine"];

const INTENT_DATA = ["engine", "language", "intent_input"];

const TTS_DATA = ["engine", "language", "voice", "tts_input"];

const STAGES: Record<PipelineRun["stage"], number> = {
  ready: 0,
  wake_word: 1,
  stt: 2,
  intent: 3,
  tts: 4,
  done: 5,
  error: 6,
};

const hasStage = (run: PipelineRun, stage: PipelineRun["stage"]) =>
  run.init_options
    ? STAGES[run.init_options.start_stage] <= STAGES[stage] &&
      STAGES[stage] <= STAGES[run.init_options.end_stage]
    : stage in run;

const maybeRenderError = (
  run: PipelineRun,
  stage: string,
  lastRunStage: string
) => {
  if (!("error" in run) || lastRunStage !== stage) {
    return "";
  }

  return html`
    <ha-alert alert-type="error">
      ${run.error!.message} (${run.error!.code})
    </ha-alert>
  `;
};

const renderProgress = (
  hass: HomeAssistant,
  pipelineRun: PipelineRun,
  stage: PipelineRun["stage"],
  start_suffix = "-start"
) => {
  const startEvent = pipelineRun.events.find(
    (ev) => ev.type === `${stage}` + start_suffix
  );
  const finishEvent = pipelineRun.events.find(
    (ev) => ev.type === `${stage}-end`
  );

  if (!startEvent) {
    return "";
  }

  if (!finishEvent) {
    if ("error" in pipelineRun) {
      return html`❌`;
    }
    return html` <ha-spinner size="small"></ha-spinner> `;
  }

  const duration =
    new Date(finishEvent.timestamp).getTime() -
    new Date(startEvent.timestamp).getTime();
  const durationString = formatNumber(duration / 1000, hass.locale, {
    maximumFractionDigits: 2,
  });
  return html`${durationString}s ✅`;
};

const renderData = (
  hass: HomeAssistant,
  data: Record<string, any>,
  keys: string[]
) =>
  keys.map((key) => {
    const label = hass.localize(
      `ui.panel.config.voice_assistants.debug.stages.${key}` as LocalizeKeys
    );
    return html`
      <div class="row">
        <div>${label}</div>
        <div>${data[key]}</div>
      </div>
    `;
  });

const dataMinusKeysRender = (
  hass: HomeAssistant,
  data: Record<string, any>,
  keys: string[]
) => {
  const result = {};
  let render = false;
  for (const key in data) {
    if (keys.includes(key) || key === "done") {
      continue;
    }
    render = true;
    result[key] = data[key];
  }
  return render
    ? html`<ha-expansion-panel class="yaml-expansion">
        <span slot="header"
          >${hass.localize("ui.panel.config.voice_assistants.debug.raw")}</span
        >
        <ha-yaml-editor readOnly autoUpdate .value=${result}></ha-yaml-editor>
      </ha-expansion-panel>`
    : "";
};

@customElement("assist-render-pipeline-run")
export class AssistPipelineDebug extends LitElement {
  @property({ attribute: false }) public hass!: HomeAssistant;

  @property({ attribute: false }) public pipelineRun!: PipelineRun;

  @property({ attribute: false }) public chatLog?: ChatLog;

  private _audioElement?: HTMLAudioElement;

  private get _isPlaying(): boolean {
    return this._audioElement != null && !this._audioElement.paused;
  }

  protected render(): TemplateResult {
    const lastRunStage: string = this.pipelineRun
      ? ["tts", "intent", "stt", "wake_word"].find(
          (stage) => stage in this.pipelineRun
        ) || "ready"
      : "ready";

    let messages: ChatLogContent[];

    if (this.chatLog) {
      messages = this.chatLog.content.filter(
        this.pipelineRun.finished
          ? (content: ChatLogContent) =>
              content.role === "system" ||
              (content.created >= this.pipelineRun.started &&
                content.created <= this.pipelineRun.finished!)
          : (content: ChatLogContent) =>
              content.role === "system" ||
              content.created >= this.pipelineRun.started
      );
    } else {
      messages = [];

      // We don't have the chat log everywhere yet, just fallback for now.
      const userMessage =
        (this.pipelineRun.init_options &&
        "text" in this.pipelineRun.init_options.input
          ? this.pipelineRun.init_options.input.text
          : undefined) ||
        this.pipelineRun?.stt?.stt_output?.text ||
        this.pipelineRun?.intent?.intent_input;

      if (userMessage) {
        messages.push({
          role: "user",
          content: userMessage,
        } as ChatLogUserContent);
      }

      if (
        this.pipelineRun?.intent?.intent_output?.response?.speech?.plain?.speech
      ) {
        messages.push({
          role: "assistant",
          content:
            this.pipelineRun.intent.intent_output.response.speech.plain.speech,
        } as ChatLogAssistantContent);
      }
    }

    return html`
      <ha-card>
        <div class="card-content">
          <div class="row heading">
            <div>
              ${this.hass.localize(
                "ui.panel.config.voice_assistants.debug.run"
              )}
            </div>
            <div>${this.pipelineRun.stage}</div>
          </div>

          ${renderData(this.hass, this.pipelineRun.run, RUN_DATA)}
          ${messages.length > 0
            ? html`
                <div class="messages">
                  ${messages.map((content) =>
                    content.role === "system"
                      ? content.content
                        ? html`
                            <ha-expansion-panel
                              class="content-expansion ${content.role}"
                            >
                              <div slot="header">System</div>
                              <pre>${content.content}</pre>
                            </ha-expansion-panel>
                          `
                        : nothing
                      : content.role === "tool_result"
                        ? html`
                            <ha-expansion-panel
                              class="content-expansion ${content.role}"
                            >
                              <div slot="header">
                                Result for ${content.tool_name}
                              </div>
                              <ha-yaml-editor
                                read-only
                                auto-update
                                .value=${content}
                              ></ha-yaml-editor>
                            </ha-expansion-panel>
                          `
                        : html`
                            ${content.content
                              ? html`
                                  <div class=${`message ${content.role}`}>
                                    ${content.content}
                                  </div>
                                `
                              : nothing}
                            ${content.role === "assistant" &&
                            content.tool_calls?.length
                              ? html`
                                  <ha-expansion-panel
                                    class="content-expansion assistant"
                                  >
                                    <span slot="header">
                                      Call
                                      ${content.tool_calls.length === 1
                                        ? content.tool_calls[0].tool_name
                                        : `${content.tool_calls.length} tools`}
                                    </span>

                                    <ha-yaml-editor
                                      read-only
                                      auto-update
                                      .value=${content.tool_calls}
                                    ></ha-yaml-editor>
                                  </ha-expansion-panel>
                                `
                              : nothing}
                          `
                  )}
                </div>
                <div style="clear:both"></div>
              `
            : ""}
        </div>
      </ha-card>

      ${maybeRenderError(this.pipelineRun, "ready", lastRunStage)}
      ${hasStage(this.pipelineRun, "wake_word")
        ? html`
            <ha-card>
              <div class="card-content">
                <div class="row heading">
                  <span
                    >${this.hass.localize(
                      "ui.panel.config.voice_assistants.debug.stages.wake_word"
                    )}</span
                  >
                  ${renderProgress(this.hass, this.pipelineRun, "wake_word")}
                </div>
                ${this.pipelineRun.wake_word
                  ? html`
                      <div class="card-content">
                        ${renderData(
                          this.hass,
                          this.pipelineRun.wake_word,
                          WAKE_WORD_DATA
                        )}
                        ${this.pipelineRun.wake_word.wake_word_output
                          ? html`<div class="row">
                                <div>
                                  ${this.hass.localize(
                                    "ui.panel.config.voice_assistants.debug.stages.model"
                                  )}
                                </div>
                                <div>
                                  ${this.pipelineRun.wake_word.wake_word_output
                                    .ww_id}
                                </div>
                              </div>
                              <div class="row">
                                <div>
                                  ${this.hass.localize(
                                    "ui.panel.config.voice_assistants.debug.stages.timestamp"
                                  )}
                                </div>
                                <div>
                                  ${this.pipelineRun.wake_word.wake_word_output
                                    .timestamp}
                                </div>
                              </div>`
                          : ""}
                        ${dataMinusKeysRender(
                          this.hass,
                          this.pipelineRun.wake_word,
                          WAKE_WORD_DATA
                        )}
                      </div>
                    `
                  : ""}
              </div>
            </ha-card>
          `
        : ""}
      ${maybeRenderError(this.pipelineRun, "wake_word", lastRunStage)}
      ${hasStage(this.pipelineRun, "stt")
        ? html`
            <ha-card>
              <div class="card-content">
                <div class="row heading">
                  <span
                    >${this.hass.localize(
                      "ui.panel.config.voice_assistants.debug.stages.speech_to_text"
                    )}</span
                  >
                  ${renderProgress(
                    this.hass,
                    this.pipelineRun,
                    "stt",
                    "-vad-end"
                  )}
                </div>
                ${this.pipelineRun.stt
                  ? html`
                      <div class="card-content">
                        ${renderData(this.hass, this.pipelineRun.stt, STT_DATA)}
                        <div class="row">
                          <div>
                            ${this.hass.localize(
                              "ui.panel.config.voice_assistants.debug.stages.language"
                            )}
                          </div>
                          <div>${this.pipelineRun.stt.metadata.language}</div>
                        </div>
                        ${this.pipelineRun.stt.stt_output
                          ? html`<div class="row">
                              <div>
                                ${this.hass.localize(
                                  "ui.panel.config.voice_assistants.debug.stages.output"
                                )}
                              </div>
                              <div>${this.pipelineRun.stt.stt_output.text}</div>
                            </div>`
                          : ""}
                        ${dataMinusKeysRender(
                          this.hass,
                          this.pipelineRun.stt,
                          STT_DATA
                        )}
                      </div>
                    `
                  : ""}
              </div>
            </ha-card>
          `
        : ""}
      ${maybeRenderError(this.pipelineRun, "stt", lastRunStage)}
      ${hasStage(this.pipelineRun, "intent")
        ? html`
            <ha-card>
              <div class="card-content">
                <div class="row heading">
                  <span
                    >${this.hass.localize(
                      "ui.panel.config.voice_assistants.debug.stages.natural_language_processing"
                    )}</span
                  >
                  ${renderProgress(this.hass, this.pipelineRun, "intent")}
                </div>
                ${this.pipelineRun.intent
                  ? html`
                      <div class="card-content">
                        ${renderData(
                          this.hass,
                          this.pipelineRun.intent,
                          INTENT_DATA
                        )}
                        ${this.pipelineRun.intent.intent_output
                          ? html`<div class="row">
                                <div>
                                  ${this.hass.localize(
                                    "ui.panel.config.voice_assistants.debug.stages.response_type"
                                  )}
                                </div>
                                <div>
                                  ${this.pipelineRun.intent.intent_output
                                    .response.response_type}
                                </div>
                              </div>
                              ${this.pipelineRun.intent.intent_output.response
                                .response_type === "error"
                                ? html`<div class="row">
                                    <div>
                                      ${this.hass.localize(
                                        "ui.panel.config.voice_assistants.debug.error.code"
                                      )}
                                    </div>
                                    <div>
                                      ${this.pipelineRun.intent.intent_output
                                        .response.data.code}
                                    </div>
                                  </div>`
                                : ""}`
                          : ""}
                        <div class="row">
                          <div>
                            ${this.hass.localize(
                              "ui.panel.config.voice_assistants.debug.stages.prefer_local"
                            )}
                          </div>
                          <div>
                            ${this.pipelineRun.intent.prefer_local_intents}
                          </div>
                        </div>
                        <div class="row">
                          <div>
                            ${this.hass.localize(
                              "ui.panel.config.voice_assistants.debug.stages.processed_locally"
                            )}
                          </div>
                          <div>
                            ${this.pipelineRun.intent.processed_locally}
                          </div>
                        </div>
                        ${dataMinusKeysRender(
                          this.hass,
                          this.pipelineRun.intent,
                          INTENT_DATA
                        )}
                      </div>
                    `
                  : ""}
              </div>
            </ha-card>
          `
        : ""}
      ${maybeRenderError(this.pipelineRun, "intent", lastRunStage)}
      ${hasStage(this.pipelineRun, "tts")
        ? html`
            <ha-card>
              <div class="card-content">
                <div class="row heading">
                  <span
                    >${this.hass.localize(
                      "ui.panel.config.voice_assistants.debug.stages.text_to_speech"
                    )}</span
                  >
                  ${renderProgress(this.hass, this.pipelineRun, "tts")}
                </div>
                ${this.pipelineRun.tts
                  ? html`
                      <div class="card-content">
                        ${renderData(this.hass, this.pipelineRun.tts, TTS_DATA)}
                        ${dataMinusKeysRender(
                          this.hass,
                          this.pipelineRun.tts,
                          TTS_DATA
                        )}
                      </div>
                    `
                  : ""}
              </div>
              ${this.pipelineRun?.tts?.tts_output
                ? html`
                    <div class="card-actions">
                      <ha-button
                        .variant=${this._isPlaying ? "danger" : "brand"}
                        @click=${this._isPlaying
                          ? this._stopTTS
                          : this._playTTS}
                      >
                        ${this._isPlaying
                          ? this.hass.localize(
                              "ui.panel.config.voice_assistants.debug.stop_audio"
                            )
                          : this.hass.localize(
                              "ui.panel.config.voice_assistants.debug.play_audio"
                            )}
                      </ha-button>
                    </div>
                  `
                : ""}
            </ha-card>
          `
        : ""}
      ${maybeRenderError(this.pipelineRun, "tts", lastRunStage)}
      <ha-card>
        <ha-expansion-panel class="yaml-expansion">
          <span slot="header"
            >${this.hass.localize(
              "ui.panel.config.voice_assistants.debug.raw"
            )}</span
          >
          <ha-yaml-editor
            read-only
            auto-update
            .value=${this.pipelineRun}
          ></ha-yaml-editor>
        </ha-expansion-panel>
      </ha-card>
    `;
  }

  private _playTTS(): void {
    // Stop any existing audio first
    this._stopTTS();

    const url = this.pipelineRun!.tts!.tts_output!.url;
    this._audioElement = new Audio(url);

    this._audioElement.addEventListener("error", () => {
      showAlertDialog(this, {
        title: this.hass.localize(
          "ui.panel.config.voice_assistants.debug.error.title"
        ),
        text: this.hass.localize(
          "ui.panel.config.voice_assistants.debug.error.playing_audio"
        ),
      });
    });

    this._audioElement.addEventListener("play", () => {
      this.requestUpdate();
    });

    this._audioElement.addEventListener("ended", () => {
      this.requestUpdate();
    });

    this._audioElement.addEventListener("canplaythrough", () => {
      this._audioElement!.play();
    });
  }

  private _stopTTS(): void {
    if (this._audioElement) {
      this._audioElement.pause();
      this._audioElement.currentTime = 0;
      this._audioElement = undefined;
      this.requestUpdate();
    }
  }

  public disconnectedCallback(): void {
    super.disconnectedCallback();
    this._stopTTS();
  }

  static styles = css`
    :host {
      display: block;
    }
    ha-card,
    ha-alert {
      display: block;
      margin-bottom: 16px;
    }
    .row {
      display: flex;
      justify-content: space-between;
    }
    .row > div:last-child {
      text-align: right;
    }
    .yaml-expansion {
      padding-left: 8px;
      padding-inline-start: 8px;
      padding-inline-end: initial;
    }
    .card-content .yaml-expansion {
      padding-left: 0px;
      padding-inline-start: 0px;
      padding-inline-end: initial;
      --expansion-panel-summary-padding: 0px;
      --expansion-panel-content-padding: 0px;
    }
    .heading {
      font-weight: var(--ha-font-weight-medium);
      margin-bottom: 16px;
    }

    .messages {
      margin-top: 8px;
    }

    .content-expansion {
      margin: 8px 0;
      border-radius: var(--ha-border-radius-xl);
      clear: both;
      padding: 0 8px;
      --input-fill-color: none;
      max-width: calc(100% - 24px);
      --expansion-panel-summary-padding: 0px;
      --expansion-panel-content-padding: 0px;
    }

    .content-expansion *[slot="header"] {
      font-weight: var(--ha-font-weight-normal);
    }

    .system {
      background-color: var(--success-color);
    }

    .message {
      padding: 8px;
    }

    .message,
    .content-expansion {
      font-size: var(--ha-font-size-l);
      margin: 8px 0;
      border-radius: var(--ha-border-radius-xl);
      clear: both;
    }

    .messages pre {
      white-space: pre-wrap;
    }

    .user,
    .tool_result {
      margin-left: 24px;
      margin-inline-start: 24px;
      margin-inline-end: initial;
      float: var(--float-end);
      border-bottom-right-radius: 0px;
      background-color: var(--light-primary-color);
      color: var(--text-light-primary-color, var(--primary-text-color));
      direction: var(--direction);
    }

    .tool_result [slot="header"] {
      color: var(--text-light-primary-color, var(--primary-text-color));
    }

    .message.user,
    .content-expansion div[slot="header"] {
      text-align: right;
    }

    .assistant {
      margin-right: 24px;
      margin-inline-end: 24px;
      margin-inline-start: initial;
      float: var(--float-start);
      border-bottom-left-radius: 0px;
      background-color: var(--primary-color);
      color: var(--text-primary-color);
      direction: var(--direction);
    }
  `;
}

declare global {
  interface HTMLElementTagNameMap {
    "assist-render-pipeline-run": AssistPipelineDebug;
  }
}
