import { useMemo, useState } from 'react';
import { Link } from 'react-router-dom';
import type { BindingPartner, Member, PartnerName, StateInfo } from '../lib/api';
import { chainPath, fmtInt } from '../lib/format';
import { EmptyState, ExternalLink } from './ui';

/** Key of an observation's biological binding partners that are not copies of itself. */
export const SELF = 'self';
export const NO_PARTNER = '__none__';

const TYPE_LABEL: Record<string, string> = { protein: 'Protein', dna: 'DNA', rna: 'RNA', hybrid: 'DNA/RNA', other: 'Polymer' };

/** Display label for a partner key ('self', a UniProt accession, 'dna', 'protein:<description>'). */
export function partnerLabel(key: string, names?: Record<string, PartnerName>): string {
  if (key === SELF) return 'Same protein';
  if (key === NO_PARTNER) return 'No partner';
  const n = names?.[key];
  if (n?.description) return n.description;
  const i = key.indexOf(':');
  if (i > 0) return key.slice(i + 1);
  return TYPE_LABEL[key] ?? key;
}

export function partnerType(key: string, names?: Record<string, PartnerName>): string {
  if (key === SELF) return 'self';
  return names?.[key]?.type ?? (key.includes(':') ? key.split(':')[0] : TYPE_LABEL[key] ? key : 'protein');
}

/** Readable name for a partner chain; oligonucleotides described by their sequence become "DNA strand (21 nt)". */
export function polymerName(description: string | null, type: string, entity?: string | null): string {
  if (!description) return entity ? `Entity ${entity}` : TYPE_LABEL[type] ?? type;
  if ((type === 'dna' || type === 'rna' || type === 'hybrid') && /\*|^5'|\(5'/.test(description)) {
    const n = (description.match(/\*/g) ?? []).length;
    return `${TYPE_LABEL[type]} strand${n ? ` (${n} nt)` : ''}`;
  }
  return description;
}

function short(s: string, max = 28) {
  return s.length > max ? `${s.slice(0, max - 1)}…` : s;
}

export function sameSet(a: string[] = [], b: string[] = []) {
  if (a.length !== b.length) return false;
  const s = new Set(a);
  return b.every((k) => s.has(k));
}

/** Compact partner chips; with `home`, partners absent from it are marked as gained and missing ones as lost. */
export function PartnerChips({
  keys = [],
  names,
  home,
  max = 2,
  stacked = false,
}: {
  keys?: string[];
  names?: Record<string, PartnerName>;
  home?: string[];
  max?: number;
  /** one chip per line with the full name (for narrow table cells) */
  stacked?: boolean;
}) {
  const homeSet = home ? new Set(home) : null;
  const lost = home ? home.filter((k) => !keys.includes(k)) : [];
  const title = [
    keys.length ? `Partners: ${keys.map((k) => partnerLabel(k, names)).join(', ')}` : 'No biological polymer partner',
    lost.length ? `Not bound here (bound in this chain): ${lost.map((k) => partnerLabel(k, names)).join(', ')}` : '',
  ]
    .filter(Boolean)
    .join('\n');
  if (!keys.length && !lost.length) return <span className="muted">—</span>;
  return (
    <span className={`pchips ${stacked ? 'pchips--stacked' : ''}`} title={title}>
      {keys.slice(0, max).map((k) => (
        <span
          key={k}
          className={`pchip pchip--${partnerType(k, names)} ${homeSet && !homeSet.has(k) ? 'is-gained' : ''}`}
        >
          {stacked ? partnerLabel(k, names) : short(partnerLabel(k, names), 22)}
        </span>
      ))}
      {keys.length > max && <span className="pchip pchip--more">+{keys.length - max}</span>}
      {!keys.length && <span className="muted small">none</span>}
      {lost.length > 0 && <span className="pchip pchip--lost">−{lost.length}</span>}
    </span>
  );
}

/* ───────────────────────────── chain page: every contacting polymer chain */

function kindOf(p: BindingPartner): { label: string; tone: string; hint: string } {
  if (p.biological)
    return p.same_entity
      ? { label: 'Same protein', tone: 'homo', hint: "Another copy of this chain's protein (same entity), in contact within biological assembly 1" }
      : { label: 'Bound partner', tone: 'bio', hint: 'A different polymer, in contact within biological assembly 1' };
  if (p.n_res_target < 5)
    return { label: 'Marginal', tone: 'weak', hint: 'Fewer than 5 interface residues on this chain' };
  return {
    label: 'Crystal contact',
    tone: 'weak',
    hint: p.via_symmetry
      ? 'Touches only a symmetry mate and is not part of biological assembly 1 (crystal packing)'
      : 'In the deposited coordinates but not in biological assembly 1 (likely crystal packing)',
  };
}

function InterfaceStrip({ ranges, length }: { ranges: Array<[number, number]>; length: number }) {
  if (!length) return null;
  return (
    <span className="istrip" role="img" aria-label={`Interface residues ${ranges.map(([a, b]) => (a === b ? a : `${a}-${b}`)).join(', ')}`}>
      {ranges.map(([a, b]) => (
        <span
          key={a}
          className="istrip__seg"
          style={{ left: `${((a - 1) / length) * 100}%`, width: `max(2px, ${((b - a + 1) / length) * 100}%)` }}
        />
      ))}
    </span>
  );
}

function fmtRanges(r: Array<[number, number]>) {
  return r.map(([a, b]) => (a === b ? `${a}` : `${a}–${b}`)).join(', ');
}

export function BindingPartners({
  partners,
  pdb,
  length,
}: {
  partners: BindingPartner[];
  pdb: string;
  length: number;
}) {
  const nWeak = partners.filter((p) => !p.biological).length;
  const [showAll, setShowAll] = useState(false);
  const rows = showAll ? partners : partners.filter((p) => p.biological);

  if (!partners.length)
    return <EmptyState title="No polymer contacts">This chain touches no other protein, DNA or RNA chain within 4.5 Å.</EmptyState>;

  return (
    <>
      {rows.length === 0 ? (
        <p className="muted">No biological partner: every contact is a crystal contact or marginal.</p>
      ) : (
        <div className="table-wrap">
          <table className="table table--dense">
            <thead>
              <tr>
                <th>Partner</th>
                <th>Chain</th>
                <th>Kind</th>
                <th className="num" title="Interface residues on this chain / on the partner">
                  Interface
                </th>
                <th>Interface on this chain (SEQRES)</th>
              </tr>
            </thead>
            <tbody>
              {rows.map((p) => {
                const k = kindOf(p);
                return (
                  <tr key={p.chain} className={p.biological ? '' : 'is-dim'}>
                    <td>
                      <span className="partner-name">
                        <span className={`ptype ptype--${p.type}`}>{TYPE_LABEL[p.type] ?? p.type}</span>
                        <span className="partner-name__text" title={p.description ?? undefined}>
                          {polymerName(p.description, p.type, p.entity)}
                        </span>
                      </span>
                      {p.uniprot && (
                        <span className="partner-name__ids">
                          {p.uniprot.split(';').map((u) => (
                            <ExternalLink key={u} href={`https://www.uniprot.org/uniprotkb/${u}`}>
                              {u}
                            </ExternalLink>
                          ))}
                        </span>
                      )}
                    </td>
                    <td className="mono">
                      {p.observation ? (
                        <Link to={chainPath(pdb, p.chain)} title={`${p.observation.sequence_id}, state ${p.observation.state_id ?? '?'}`}>
                          {pdb.toUpperCase()}.{p.chain}
                        </Link>
                      ) : (
                        `${pdb.toUpperCase()}.${p.chain}`
                      )}
                      {p.via_symmetry && (
                        <span className="muted small" title="Contact is with a crystallographic symmetry copy of this chain">
                          {' '}
                          (sym)
                        </span>
                      )}
                    </td>
                    <td>
                      <span className={`pkind pkind--${k.tone}`} title={k.hint}>
                        {k.label}
                      </span>
                    </td>
                    <td className="num tabular" title={`${fmtInt(p.n_contacts)} atom contacts within 4.5 Å`}>
                      {p.n_res_target} / {p.n_res_partner ?? '—'}
                    </td>
                    <td title={fmtRanges(p.interface)}>
                      <InterfaceStrip ranges={p.interface} length={length} />
                    </td>
                  </tr>
                );
              })}
            </tbody>
          </table>
        </div>
      )}
      {nWeak > 0 && (
        <button type="button" className="btn btn--ghost btn--sm partners-toggle" onClick={() => setShowAll((s) => !s)}>
          {showAll ? 'Hide' : 'Show'} {fmtInt(nWeak)} crystal / marginal contact{nWeak === 1 ? '' : 's'}
        </button>
      )}
    </>
  );
}

/* ───────────────────────────── sequence page: partners × states */

const MAX_ROWS = 12;
const MAX_COLS = 14;

export function PartnersByState({
  members,
  states,
  names,
  selectedState,
  onSelectState,
  selectedPartner,
  onSelectPartner,
}: {
  members: Member[];
  states: StateInfo[];
  names?: Record<string, PartnerName>;
  selectedState: number | null;
  onSelectState: (s: number | null) => void;
  selectedPartner: string | null;
  onSelectPartner: (k: string | null) => void;
}) {
  const { rows, cols, cell, perState } = useMemo(() => {
    const total = new Map<string, number>();
    const cell = new Map<string, number>();
    const perState = new Map<number, number>();
    for (const m of members) {
      const keys = m.partners?.length ? m.partners : [NO_PARTNER];
      perState.set(m.state_id, (perState.get(m.state_id) ?? 0) + 1);
      for (const k of keys) {
        total.set(k, (total.get(k) ?? 0) + 1);
        cell.set(`${k}|${m.state_id}`, (cell.get(`${k}|${m.state_id}`) ?? 0) + 1);
      }
    }
    const rows = [...total.entries()].sort((a, b) => b[1] - a[1]).slice(0, MAX_ROWS);
    const cols = states.slice(0, MAX_COLS);
    return { rows, cols, cell, perState };
  }, [members, states]);

  if (rows.length === 1 && rows[0][0] === NO_PARTNER)
    return <p className="muted">No observation of this sequence has a biological polymer partner.</p>;

  return (
    <div className="pmatrix-wrap">
      <table className="pmatrix">
        <thead>
          <tr>
            <th className="pmatrix__corner">Partner \ state</th>
            {cols.map((s) => (
              <th key={s.state_id}>
                <button
                  type="button"
                  className={`pmatrix__state ${selectedState === s.state_id ? 'is-on' : ''}`}
                  onClick={() => onSelectState(selectedState === s.state_id ? null : s.state_id)}
                  title={`State ${s.state_id}: ${fmtInt(perState.get(s.state_id))} observations`}
                >
                  {s.state_id}
                </button>
              </th>
            ))}
            <th className="num">All</th>
          </tr>
        </thead>
        <tbody>
          {rows.map(([k, n]) => (
            <tr key={k}>
              <th scope="row">
                <button
                  type="button"
                  className={`pmatrix__partner ${selectedPartner === k ? 'is-on' : ''}`}
                  onClick={() => onSelectPartner(selectedPartner === k ? null : k)}
                  title={`${partnerLabel(k, names)}${names?.[k]?.uniprot ? ` (${names[k].uniprot})` : ''} — click to list these observations`}
                >
                  <span className={`ptype-dot ptype-dot--${partnerType(k, names)}`} aria-hidden />
                  <span className="pmatrix__label">{partnerLabel(k, names)}</span>
                  {names?.[k]?.uniprot && <span className="pmatrix__acc mono">{names[k].uniprot}</span>}
                </button>
              </th>
              {cols.map((s) => {
                const c = cell.get(`${k}|${s.state_id}`) ?? 0;
                const frac = c / (perState.get(s.state_id) || 1);
                const level = c === 0 ? 0 : Math.min(5, 1 + Math.floor(frac * 5));
                return (
                  <td
                    key={s.state_id}
                    className={`pmatrix__cell seq-bg-${level}`}
                    title={`${partnerLabel(k, names)} in state ${s.state_id}: ${c} of ${perState.get(s.state_id)} observations`}
                  >
                    {c > 0 ? c : ''}
                  </td>
                );
              })}
              <td className="num tabular">{fmtInt(n)}</td>
            </tr>
          ))}
        </tbody>
      </table>
      {(states.length > MAX_COLS) && (
        <p className="muted small">Showing the {MAX_COLS} largest of {states.length} states.</p>
      )}
    </div>
  );
}