import * as DG from 'datagrok-api/dg';

import {TAGS as bioTAGS} from '@datagrok-libraries/bio/src/utils/macromolecule';
import {
  SeqAnnotation, SeqAnnotationHit, RowAnnotationData,
  AnnotationCategory, AnnotationVisualType,
} from '@datagrok-libraries/bio/src/utils/macromolecule/annotations';
import {SeqRegion} from '../get-region-func-editor';

/** Prefix for hidden companion annotation columns (~ hides them in Datagrok). */
const ANNOTATION_COL_PREFIX = '~';

/** Reads column-level annotations from the `.annotations` tag.
 *  Falls back to `.regions` for backward compatibility. */
export function getColumnAnnotations(col: DG.Column<string>): SeqAnnotation[] {
  const annotationsTag = col.getTag(bioTAGS.annotations);
  if (annotationsTag) {
    try {
      return JSON.parse(annotationsTag) as SeqAnnotation[];
    } catch { /* fall through */ }
  }
  // Backward compat: convert legacy .regions to SeqAnnotation[]
  const regionsTag = col.getTag(bioTAGS.regions);
  if (regionsTag) {
    try {
      const regions: SeqRegion[] = JSON.parse(regionsTag);
      return regions.map((r, i) => ({
        id: `legacy-region-${i}`,
        name: r.name,
        description: r.description,
        start: r.start,
        end: r.end,
        visualType: AnnotationVisualType.Region,
        category: AnnotationCategory.Structure,
        autoGenerated: true,
      }));
    } catch { /* ignore */ }
  }
  return [];
}

/** Writes column-level annotations to the `.annotations` tag.
 *  Also keeps `.regions` in sync for backward compatibility with GetRegionFuncEditor. */
export function setColumnAnnotations(col: DG.Column<string>, annotations: SeqAnnotation[]): void {
  col.setTag(bioTAGS.annotations, JSON.stringify(annotations));
  // Keep .regions in sync with structure annotations
  const structureAnnotations = annotations.filter((a) => a.category === AnnotationCategory.Structure);
  if (structureAnnotations.length > 0) {
    const regions: SeqRegion[] = structureAnnotations
      .filter((a) => a.start != null && a.end != null)
      .map((a) => ({
        name: a.name,
        description: a.description ?? '',
        start: a.start!,
        end: a.end!,
      }));
    col.setTag(bioTAGS.regions, JSON.stringify(regions));
  }
}

/** Returns the name for the hidden companion annotation column. */
export function getAnnotationColumnName(seqColName: string): string {
  return `${ANNOTATION_COL_PREFIX}${seqColName}_annotations`;
}

/** Gets or creates the hidden companion column for per-row annotation hits. */
export function getOrCreateAnnotationColumn(df: DG.DataFrame, seqCol: DG.Column<string>): DG.Column<string> {
  const colName = getAnnotationColumnName(seqCol.name);
  let col = df.columns.byName(colName);
  if (!col) {
    col = df.columns.addNewString(colName);
    seqCol.setTag(bioTAGS.annotationColumnName, colName);
  }
  return col as DG.Column<string>;
}

/** Reads per-row annotation hits from the companion column. Uses version-based caching. */
const _rowDataCache = new WeakMap<DG.Column, {version: number; data:(RowAnnotationData | null)[]}>();

export function getRowAnnotations(annotCol: DG.Column<string>, rowIdx: number): RowAnnotationData | null {
  const cached = _rowDataCache.get(annotCol);
  if (cached && cached.version === annotCol.version) {
    if (cached.data[rowIdx] !== undefined)
      return cached.data[rowIdx];
  }
  // Parse this row
  const raw = annotCol.get(rowIdx);
  if (!raw) return null;
  try {
    return JSON.parse(raw) as RowAnnotationData;
  } catch { return null; }
}

/** Parses and caches all row annotations for the column. Call once when version changes. */
export function cacheAllRowAnnotations(annotCol: DG.Column<string>): (RowAnnotationData | null)[] {
  const cached = _rowDataCache.get(annotCol);
  if (cached && cached.version === annotCol.version)
    return cached.data;

  const data: (RowAnnotationData | null)[] = new Array(annotCol.length);
  for (let i = 0; i < annotCol.length; i++) {
    const raw = annotCol.get(i);
    if (!raw) {
      data[i] = null;
      continue;
    }
    try {
      data[i] = JSON.parse(raw) as RowAnnotationData;
    } catch {
      data[i] = null;
    }
  }
  _rowDataCache.set(annotCol, {version: annotCol.version, data});
  return data;
}

/** Writes per-row annotation hits. */
export function setRowAnnotations(annotCol: DG.Column<string>, rowIdx: number, hits: SeqAnnotationHit[]): void {
  annotCol.set(rowIdx, hits.length > 0 ? JSON.stringify(hits) : '');
}

/** Clears all annotations from a column (both column-level and row-level). */
export function clearAnnotations(df: DG.DataFrame, seqCol: DG.Column<string>): void {
  seqCol.setTag(bioTAGS.annotations, '');
  const annotColName = getAnnotationColumnName(seqCol.name);
  const annotCol = df.columns.byName(annotColName);
  if (annotCol)
    df.columns.remove(annotColName);
  // Clear .regions too
  seqCol.setTag(bioTAGS.regions, '');
}

/** Adds an annotation to the column-level list. */
export function addColumnAnnotation(col: DG.Column<string>, annotation: SeqAnnotation): void {
  const existing = getColumnAnnotations(col);
  existing.push(annotation);
  setColumnAnnotations(col, existing);
}

/** Removes an annotation by id from the column-level list. */
export function removeColumnAnnotation(col: DG.Column<string>, annotationId: string): void {
  const existing = getColumnAnnotations(col).filter((a) => a.id !== annotationId);
  setColumnAnnotations(col, existing);
}

/** Merges row-level annotation hits by replacing hits of one kind while preserving the rest.
 *  @param existingHits Current per-row hits
 *  @param newHits New hits to add
 *  @param replaceRegions If true, removes existing region span hits (endPositionIndex set) before merging.
 *  @param replaceLiabilities If true, removes existing non-region hits before merging. */
export function mergeRowHits(
  existingHits: SeqAnnotationHit[],
  newHits: SeqAnnotationHit[],
  replaceRegions: boolean,
  replaceLiabilities: boolean,
): SeqAnnotationHit[] {
  let kept = existingHits;
  if (replaceRegions)
    kept = kept.filter((h) => h.endPositionIndex == null);
  if (replaceLiabilities)
    kept = kept.filter((h) => h.endPositionIndex != null);
  return [...kept, ...newHits];
}