import { Task } from "@lit/task";
import { LitElement, type PropertyValueMap, css, html } from "lit";
import { customElement, property } from "lit/decorators.js";
import type { GetISOBMFFFileTranscriptionResult } from "../../../api/src/index.js";
import { EF_INTERACTIVE } from "../EF_INTERACTIVE.js";
import { CrossUpdateController } from "./CrossUpdateController.js";
import { EFAudio } from "./EFAudio.js";
import { EFSourceMixin } from "./EFSourceMixin.js";
import { EFTemporal } from "./EFTemporal.js";
import { EFVideo } from "./EFVideo.js";
import { FetchMixin } from "./FetchMixin.js";

interface WordSegment {
  text: string;
  start: number;
  end: number;
}

interface Segment {
  start: number;
  end: number;
  text: string;
}

interface Caption {
  segments: Segment[];
  word_segments: WordSegment[];
}

const stopWords = new Set(["", ".", "!", "?", ","]);

@customElement("ef-captions-active-word")
export class EFCaptionsActiveWord extends EFTemporal(LitElement) {
  static styles = [
    css`
      :host {
        display: inline-block;
        white-space: pre;
      }
      :host([hidden]) {
        display: none;
      }
    `,
  ];

  render() {
    if (stopWords.has(this.wordText)) {
      this.hidden = true;
      return undefined;
    }
    this.hidden = false;
    return html` ${this.wordText.trim()} `;
  }

  @property({ type: Number, attribute: false })
  wordStartMs = 0;

  @property({ type: Number, attribute: false })
  wordEndMs = 0;

  @property({ type: String, attribute: false })
  wordText = "";

  @property({ type: Boolean, reflect: true })
  hidden = false;

  get startTimeMs() {
    return this.wordStartMs || 0;
  }

  get durationMs(): number {
    return this.wordEndMs - this.wordStartMs;
  }
}

@customElement("ef-captions-segment")
export class EFCaptionsSegment extends EFTemporal(LitElement) {
  static styles = [
    css`
      :host {
        display: block;
      }
      :host([hidden]) {
        display: none;
      }
    `,
  ];

  render() {
    if (stopWords.has(this.segmentText)) {
      this.hidden = true;
      return undefined;
    }
    this.hidden = false;
    return html`${this.segmentText}`;
  }

  @property({ type: Number, attribute: false })
  segmentStartMs = 0;

  @property({ type: Number, attribute: false })
  segmentEndMs = 0;

  @property({ type: String, attribute: false })
  segmentText = "";

  @property({ type: Boolean, reflect: true })
  hidden = false;

  get startTimeMs() {
    return this.segmentStartMs || 0;
  }

  get durationMs(): number {
    return this.segmentEndMs - this.segmentStartMs;
  }
}

@customElement("ef-captions-before-active-word")
export class EFCaptionsBeforeActiveWord extends EFCaptionsSegment {
  static styles = [
    css`
      :host {
        display: inline-block;
        white-space: pre;
      }
      :host([hidden]) {
        display: none;
      }
    `,
  ];

  render() {
    if (stopWords.has(this.segmentText)) {
      this.hidden = true;
      return undefined;
    }
    this.hidden = false;
    return html` ${this.segmentText}`;
  }

  @property({ type: Boolean, reflect: true })
  hidden = false;

  @property({ type: String, attribute: false })
  segmentText = "";

  @property({ type: Number, attribute: false })
  segmentStartMs = 0;

  @property({ type: Number, attribute: false })
  segmentEndMs = 0;

  get startTimeMs() {
    return this.segmentStartMs || 0;
  }

  get durationMs(): number {
    return this.segmentEndMs - this.segmentStartMs;
  }
}

@customElement("ef-captions-after-active-word")
export class EFCaptionsAfterActiveWord extends EFCaptionsSegment {
  static styles = [
    css`
      :host {
        display: inline-block;
        white-space: pre;
      }
      :host([hidden]) {
        display: none;
      }
    `,
  ];

  render() {
    if (stopWords.has(this.segmentText)) {
      this.hidden = true;
      return undefined;
    }
    this.hidden = false;
    return html`${this.segmentText} `;
  }

  @property({ type: Boolean, reflect: true })
  hidden = false;

  @property({ type: String, attribute: false })
  segmentText = "";

  @property({ type: Number, attribute: false })
  segmentStartMs = 0;

  @property({ type: Number, attribute: false })
  segmentEndMs = 0;

  get startTimeMs() {
    return this.segmentStartMs || 0;
  }

  get durationMs(): number {
    return this.segmentEndMs - this.segmentStartMs;
  }
}

@customElement("ef-captions")
export class EFCaptions extends EFSourceMixin(
  EFTemporal(FetchMixin(LitElement)),
  { assetType: "caption_files" },
) {
  static styles = [
    css`
      :host {
        display: flex;
        flex-wrap: wrap;
        align-items: baseline;
        width: fit-content;
      }
      ::slotted(*) {
        margin: 0;
        padding: 0;
      }
    `,
  ];

  @property({ type: String, attribute: "display-mode", reflect: true })
  displayMode: "word" | "segment" | "context" = "segment";

  @property({ type: Number, attribute: "context-words", reflect: true })
  contextWords = 3;

  @property({ type: String, attribute: "target", reflect: true })
  targetSelector = "";

  set target(value: string) {
    this.targetSelector = value;
  }

  @property({ attribute: "word-style" })
  wordStyle = "";

  activeWordContainers = this.getElementsByTagName("ef-captions-active-word");
  segmentContainers = this.getElementsByTagName("ef-captions-segment");
  beforeActiveWordContainers = this.getElementsByTagName(
    "ef-captions-before-active-word",
  );
  afterActiveWordContainers = this.getElementsByTagName(
    "ef-captions-after-active-word",
  );

  render() {
    return html`<slot></slot>`;
  }

  transcriptionsPath() {
    if (this.targetElement.assetId) {
      return `${this.apiHost}/api/v1/isobmff_files/${this.targetElement.assetId}/transcription`;
    }
    return null;
  }

  captionsPath() {
    if (this.targetElement.assetId) {
      return `${this.apiHost}/api/v1/caption_files/${this.targetElement.assetId}`;
    }
    const targetSrc = this.targetElement.src;
    return `/@ef-captions/${targetSrc}`;
  }

  protected md5SumLoader = new Task(this, {
    autoRun: false,
    args: () => [this.target, this.fetch] as const,
    task: async ([_target, fetch], { signal }) => {
      const md5Path = `/@ef-asset/${this.targetElement.src ?? ""}`;
      const response = await fetch(md5Path, { method: "HEAD", signal });
      return response.headers.get("etag") ?? undefined;
    },
  });

  private transcriptionDataTask = new Task(this, {
    autoRun: EF_INTERACTIVE,
    args: () => [this.transcriptionsPath(), this.fetch] as const,
    task: async ([transcriptionsPath, fetch], { signal }) => {
      if (!transcriptionsPath) {
        return null;
      }
      const response = await fetch(transcriptionsPath, { signal });
      return response.json() as any as GetISOBMFFFileTranscriptionResult;
    },
  });

  private transcriptionFragmentPath(
    transcriptionId: string,
    fragmentIndex: number,
  ) {
    return `${this.apiHost}/api/v1/transcriptions/${transcriptionId}/fragments/${fragmentIndex}`;
  }

  private fragmentIndexTask = new Task(this, {
    autoRun: EF_INTERACTIVE,
    args: () =>
      [this.transcriptionDataTask.value, this.ownCurrentTimeMs] as const,
    task: async ([transcription, ownCurrentTimeMs]) => {
      if (!transcription) {
        return null;
      }
      const fragmentIndex = Math.floor(
        ownCurrentTimeMs / transcription.work_slice_ms,
      );
      return fragmentIndex;
    },
  });

  private transcriptionFragmentDataTask = new Task(this, {
    autoRun: EF_INTERACTIVE,
    args: () =>
      [
        this.transcriptionDataTask.value,
        this.fragmentIndexTask.value,
        this.fetch,
      ] as const,
    task: async ([transcription, fragmentIndex, fetch], { signal }) => {
      if (
        transcription === null ||
        transcription === undefined ||
        fragmentIndex === null ||
        fragmentIndex === undefined
      ) {
        return null;
      }
      const fragmentPath = this.transcriptionFragmentPath(
        transcription.id,
        fragmentIndex,
      );
      const response = await fetch(fragmentPath, { signal });
      return response.json() as any as Caption;
    },
  });

  frameTask = new Task(this, {
    autoRun: EF_INTERACTIVE,
    args: () => [this.transcriptionFragmentDataTask.status],
    task: async () => {
      await this.transcriptionFragmentDataTask.taskComplete;
    },
  });

  connectedCallback() {
    super.connectedCallback();
    if (this.targetElement) {
      new CrossUpdateController(this.targetElement, this);
    }
  }

  protected updated(
    _changedProperties: PropertyValueMap<any> | Map<PropertyKey, unknown>,
  ): void {
    this.updateTextContainers();
  }

  updateTextContainers() {
    const transcriptionFragment = this.transcriptionFragmentDataTask
      .value as Caption;
    if (!transcriptionFragment) {
      return;
    }

    const currentTimeMs = this.targetElement.currentSourceTimeMs;
    const currentTimeSec = currentTimeMs / 1000;

    // Find the current word from word_segments
    const currentWord = transcriptionFragment.word_segments.find(
      (word) => currentTimeSec >= word.start && currentTimeSec <= word.end,
    );

    // Find the current segment
    const currentSegment = transcriptionFragment.segments.find(
      (segment) =>
        currentTimeSec >= segment.start && currentTimeSec <= segment.end,
    );

    for (const wordContainer of this.activeWordContainers) {
      if (currentWord) {
        wordContainer.wordText = currentWord.text;
        wordContainer.wordStartMs = currentWord.start * 1000;
        wordContainer.wordEndMs = currentWord.end * 1000;
      }
    }

    for (const segmentContainer of this.segmentContainers) {
      if (currentSegment) {
        segmentContainer.segmentText = currentSegment.text;
        segmentContainer.segmentStartMs = currentSegment.start * 1000;
        segmentContainer.segmentEndMs = currentSegment.end * 1000;
      }
    }

    // Only process context if we have both a current word and segment
    if (currentWord && currentSegment) {
      // Find all word segments that fall within the current segment's time range
      const segmentWords = transcriptionFragment.word_segments.filter(
        (word) =>
          word.start >= currentSegment.start && word.end <= currentSegment.end,
      );

      // Find the index of the current word within the segment's word segments
      const currentWordIndex = segmentWords.findIndex(
        (word) =>
          word.start === currentWord.start && word.end === currentWord.end,
      );

      if (currentWordIndex !== -1) {
        // Get words before the current word
        const beforeWords = segmentWords
          .slice(0, currentWordIndex)
          .map((w) => w.text.trim())
          .join(" ");

        // Get words after the current word
        const afterWords = segmentWords
          .slice(currentWordIndex + 1)
          .map((w) => w.text.trim())
          .join(" ");

        // Update before containers
        for (const container of this.beforeActiveWordContainers) {
          container.segmentText = beforeWords;
          container.segmentStartMs = currentSegment.start * 1000;
          container.segmentEndMs = currentWord.start * 1000;
        }

        // Update after containers
        for (const container of this.afterActiveWordContainers) {
          container.segmentText = afterWords;
          container.segmentStartMs = currentWord.end * 1000;
          container.segmentEndMs = currentSegment.end * 1000;
        }
      }
    }
  }

  get targetElement() {
    const target = document.getElementById(this.targetSelector ?? "");
    if (target instanceof EFAudio || target instanceof EFVideo) {
      return target;
    }
    throw new Error("Invalid target, must be an EFAudio or EFVideo element");
  }
}

declare global {
  interface HTMLElementTagNameMap {
    "ef-captions": EFCaptions;
    "ef-captions-active-word": EFCaptionsActiveWord;
    "ef-captions-segment": EFCaptionsSegment;
    "ef-captions-before-active-word": EFCaptionsBeforeActiveWord;
    "ef-captions-after-active-word": EFCaptionsAfterActiveWord;
  }
}
