import React, { useEffect, useMemo, useState } from 'react';
import { Box, CircularProgress, IconButton } from '@mui/material';
import ArrowBackIosNewIcon from '@mui/icons-material/ArrowBackIosNew';
import ArrowForwardIosIcon from '@mui/icons-material/ArrowForwardIos';
import { useAuth } from '../../../contexts/AuthContext';

type Artefact = { id: string; type: string; rel_path: string; created_at: string };

type DoclingJson = Record<string, unknown> & {
  pages?: Array<{
    image_base64?: string;
    image?: { uri?: string; image_base64?: string; mimetype?: string };
    width?: number;
    height?: number;
  }> | Record<string, unknown>;
  page_images?: Array<{ uri?: string; image_base64?: string }>;
  images?: Array<{ uri?: string; image_base64?: string }>;
  frontpage?: { image_base64?: string };
  cover?: { image_base64?: string };
};

type PageImagesManifest = {
  version: number;
  file_id: string;
  page_count: number;
  bucket?: string;
  base_dir?: string;
  page_images: Array<{
    page: number;
    full_image_path: string;
    thumbnail_path: string;
    full_dimensions?: { width: number; height: number };
    thumbnail_dimensions?: { width: number; height: number };
  }>
};

export const CCDoclingViewer: React.FC<{
  fileId: string;
  currentPage?: number;
  onPageChange?: (page: number) => void;
  onExtractedText?: (text: string) => void;
  onTotalPagesChange?: (total: number) => void;
  hideToolbar?: boolean;
  sectionRange?: { start: number; end: number };
}> = ({ fileId, currentPage, onPageChange, onExtractedText, onTotalPagesChange, hideToolbar, sectionRange }) => {
  const { accessToken } = useAuth();
  const [loading, setLoading] = useState(false);
  const [error, setError] = useState<string | null>(null);
  const [images, setImages] = useState<Array<{ src: string; width?: number; height?: number }>>([]);
  const [manifest, setManifest] = useState<PageImagesManifest | null>(null);
  const [pageLocal, setPageLocal] = useState<number>(1);

  const page = typeof currentPage === 'number' ? currentPage : pageLocal;

  const norm = (s: unknown): string => String(s ?? '')
    .replace(/\r\n/g, '\n')
    .replace(/\t/g, ' ')
    .replace(/[ \f\v]{2,}/g, ' ')
    .replace(/\n{3,}/g, '\n\n')
    .trim();

  const asRecord = (v: unknown): Record<string, unknown> => (v && typeof v === 'object') ? (v as Record<string, unknown>) : {};
  const asArray = (v: unknown): unknown[] => Array.isArray(v) ? v : [];

  const getText = (n: unknown): string => {
    const node = asRecord(n);
    const cands = [
      node['text'], node['orig'], node['content'], node['value'],
      node['md'], node['markdown'], node['plain_text'], node['caption'], node['title']
    ];
    const val = cands.find((v): v is string => typeof v === 'string' && v.trim().length > 0);
    return norm(val ?? '');
  };

  const collectSimpleText = (doc: unknown): string => {
    const docRec = asRecord(doc);
    const d = asRecord(asRecord(docRec['document'])['json_content'] ?? docRec['json_content'] ?? docRec);
    const parts: string[] = [];

    for (const t of asArray(d['texts'])) {
        const txt = getText(t);
        if (txt) parts.push(txt);
    }

    for (const li of asArray(d['lists'])) {
      const items = asArray(asRecord(li)['items']).map(getText).filter(Boolean) as string[];
        if (items.length) parts.push(norm(items.join('\n')));
    }

    for (const tbl of asArray(d['tables'])) {
      const data = asRecord(asRecord(tbl)['data']);
      const grid = asArray(data['grid']);
      if (grid.length) {
        const rows = grid.map((row) => `| ${asArray(row).map((c) => getText(c)).join(' | ')} |`);
          if (rows.length >= 2) {
            const firstLen = asArray(grid[0]).length;
            rows.splice(1, 0, `| ${Array(firstLen).fill('---').join(' | ')} |`);
          }
          if (rows.length) parts.push(norm(rows.join('\n')));
      } else {
        const rowsArr = asArray(data['rows']);
        if (rowsArr.length) {
          const rows: string[] = [];
          for (const r of rowsArr) {
            const rRec = asRecord(r);
            const cells = (asArray(rRec['cells']).length ? asArray(rRec['cells']) : asArray(r)).map((c) => getText(c));
            rows.push(`| ${cells.join(' | ')} |`);
          }
          if (rows.length) parts.push(norm(rows.join('\n')));
        }
      }
    }

    return norm(parts.join('\n\n'));
  };

  const extractImages = (rawDoc: unknown): Array<{ src: string; width?: number; height?: number }> => {
    const docRec = asRecord(rawDoc);
    const d = asRecord(asRecord(docRec['document'])['json_content'] ?? docRec['json_content'] ?? docRec);
    const out: Array<{ src: string; width?: number; height?: number }> = [];

    const pushUri = (uri?: string) => {
      if (!uri) return;
      if (uri.startsWith('data:')) out.push({ src: uri });
      else if (/^[A-Za-z0-9+/=]+$/.test(uri)) out.push({ src: `data:image/png;base64,${uri}` });
    };

    // Case 1: pages is an array with image_base64 or image.uri
    const pagesVal = d['pages'];
    if (Array.isArray(pagesVal)) {
      for (const p of pagesVal) {
        const pRec = asRecord(p);
        const image = asRecord(pRec['image']);
        const b64 = pRec['image_base64'] as string | undefined;
        const b64img = image['image_base64'] as string | undefined;
        const uri = image['uri'] as string | undefined;
        if (b64) out.push({ src: `data:image/png;base64,${b64}` });
        else if (b64img) out.push({ src: `data:image/png;base64,${b64img}` });
        else if (uri) pushUri(uri);
      }
    }

    // Case 2: pages is an object keyed by page number, each with image.uri
    if (!out.length && pagesVal && typeof pagesVal === 'object' && !Array.isArray(pagesVal)) {
      const pagesRec = asRecord(pagesVal);
      const keys = Object.keys(pagesRec).sort((a, b) => Number(a) - Number(b));
      for (const k of keys) {
        const pRec = asRecord(pagesRec[k]);
        const img = asRecord(pRec['image']);
        const uri = img['uri'] as string | undefined;
        const b64 = pRec['image_base64'] as string | undefined;
        if (uri) pushUri(uri);
        else if (b64) out.push({ src: `data:image/png;base64,${b64}` });
      }
    }

    // Case 3: page_images or images arrays with data URIs
    if (!out.length && Array.isArray(d['page_images'])) {
      for (const im of d['page_images'] as Array<Record<string, unknown>>) pushUri((im['uri'] as string | undefined) || (im['image_base64'] as string | undefined));
    }
    if (!out.length && Array.isArray(d['images'])) {
      for (const im of d['images'] as Array<Record<string, unknown>>) pushUri((im['uri'] as string | undefined) || (im['image_base64'] as string | undefined));
    }

    // Fallback: frontpage/cover only
    const front = asRecord(d['frontpage']);
    const cover = asRecord(d['cover']);
    const frontB64 = front['image_base64'] as string | undefined;
    const coverB64 = cover['image_base64'] as string | undefined;
    if (!out.length && (frontB64 || coverB64)) {
      const src = frontB64 || coverB64;
      if (src) out.push({ src: `data:image/png;base64,${src}` });
    }

    return out;
  };

  useEffect(() => {
    const run = async () => {
      if (!fileId) return;
      setLoading(true);
      setError(null);
      try {
        // Try page-images manifest first
        const API_BASE = import.meta.env.VITE_API_BASE || '/api';
        try {
          const mRes = await fetch(`${API_BASE}/database/files/${encodeURIComponent(fileId)}/page-images/manifest`, {
            headers: { 'Authorization': `Bearer ${accessToken || ''}` }
          });
          if (mRes.ok) {
            const m: PageImagesManifest = await mRes.json();
            setManifest(m);
            setImages([]); // we will render via manifest in viewer
            if (!currentPage) setPageLocal(1);
            return; // skip legacy docling path
          }
        } catch (e) {
          // ignore and fallback to legacy
        }

        // Legacy: Load artefacts for file to find docling JSON artefacts
        const artefactsRes = await fetch(`${import.meta.env.VITE_API_BASE || '/api'}/database/files/${encodeURIComponent(fileId)}/artefacts`, {
          headers: { 'Authorization': `Bearer ${accessToken || ''}` }
        });
        if (!artefactsRes.ok) throw new Error(await artefactsRes.text());
        const artefacts: Artefact[] = await artefactsRes.json();

        // Prefer full-file no-OCR artefact for complete page images
        const noocr = artefacts.find(a => a.type === 'docling_noocr_json');
        const frontmatter = artefacts.find(a => a.type === 'docling_frontmatter_json');
        const target = noocr || frontmatter;
        if (!target) {
          setError('No Docling artefacts found. Generate initial artefacts from the file menu.');
          setImages([]);
          return;
        }

        // Download artefact JSON via backend (service-role) to avoid RLS issues
        const jsonRes = await fetch(`${API_BASE}/database/files/${encodeURIComponent(fileId)}/artefacts/${encodeURIComponent(target.id)}/json`, {
          headers: { 'Authorization': `Bearer ${accessToken || ''}` }
        });
        if (!jsonRes.ok) throw new Error(await jsonRes.text());
        const doc: DoclingJson = await jsonRes.json();

        const imgs = extractImages(doc);
        setImages(imgs);

        if (onExtractedText) {
          const text = collectSimpleText(doc);
          onExtractedText(text);
        }
      } catch (e: unknown) {
        setError(e instanceof Error ? e.message : 'Failed to load document');
        setImages([]);
      } finally {
        setLoading(false);
      }
    };
    run();
  }, [fileId]); /* eslint-disable-line react-hooks/exhaustive-deps */

  const API_BASE = useMemo(() => import.meta.env.VITE_API_BASE || '/api', []);

  const pageProxyUrl = useMemo(() => {
    if (!manifest) return undefined;
    const idx = Math.max(0, Math.min((manifest.page_count || 1) - 1, (page || 1) - 1));
    const pg = manifest.page_images[idx];
    if (!pg) return undefined;
    const bucket = manifest.bucket || '';
    const path = pg.full_image_path;
    return `${API_BASE}/database/files/proxy?bucket=${encodeURIComponent(bucket)}&path=${encodeURIComponent(path)}`;
  }, [manifest, page, API_BASE]);

  const [pageObjectUrl, setPageObjectUrl] = useState<string | undefined>(undefined);
  const [cacheUrls] = useState<Map<number, string>>(() => new Map());

  useEffect(() => {
    let revoked: string | null = null;
    const load = async () => {
      if (!pageProxyUrl || !manifest) {
        setPageObjectUrl(undefined);
        return;
      }
      // Cache by page number to avoid repeated fetches
      const key = page;
      const cached = cacheUrls.get(key);
      if (cached) {
        setPageObjectUrl(cached);
        return;
      }
      const token = accessToken || '';
      let resp = await fetch(pageProxyUrl, { headers: { Authorization: `Bearer ${token}` } });
      if (!resp.ok && manifest) {
        // Fallback to thumbnail if the full image is not accessible yet
        const idx = Math.max(0, Math.min((manifest.page_count || 1) - 1, (page || 1) - 1));
        const pg = manifest.page_images[idx];
        if (pg) {
          const thumbUrl = `${API_BASE}/database/files/proxy?bucket=${encodeURIComponent(manifest.bucket || '')}&path=${encodeURIComponent(pg.thumbnail_path)}`;
          resp = await fetch(thumbUrl, { headers: { Authorization: `Bearer ${token}` } });
        }
      }
      if (!resp.ok) {
        setError(`Failed to load page ${page}: ${resp.status}`);
        setPageObjectUrl(undefined);
        return;
      }
      const blob = await resp.blob();
      const url = URL.createObjectURL(blob);
      cacheUrls.set(key, url);
      setPageObjectUrl(url);
      revoked = url;
    };
    load();
    return () => {
      // Do not revoke cached urls immediately; only revoke if it's a temp assignment
      // We keep the cache for navigation performance.
      if (revoked && ![...cacheUrls.values()].includes(revoked)) {
        URL.revokeObjectURL(revoked);
      }
    };
  }, [pageProxyUrl, manifest, page, cacheUrls]);

  const totalPages = manifest?.page_count || images.length || 1;

  // Inform parent about total pages when it changes
  useEffect(() => {
    if (onTotalPagesChange) onTotalPagesChange(totalPages);
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [totalPages]);

  const handlePageChange = (p: number) => {
    const clamped = Math.max(1, Math.min(totalPages, p));
    if (onPageChange) onPageChange(clamped);
    else setPageLocal(clamped);
  };

  const content = useMemo(() => {
    if (loading) return <Box sx={{ p: 2 }}><CircularProgress size={20} /></Box>;
    if (error) return <Box sx={{ p: 2, color: 'var(--color-text-2)' }}>{error}</Box>;
    // New single-page view using manifest
    if (manifest) {
      // Multi-page section view
      const start = sectionRange?.start ?? page;
      const end = sectionRange?.end ?? page;
      const pages: number[] = [];
      for (let p = start; p <= Math.min(end, totalPages); p++) pages.push(p);
      return (
        <Box sx={{ width: '100%', height: '100%', display: 'flex', flexDirection: 'column' }}>
          {!hideToolbar && (
            <Box sx={{ p: 1, display: 'flex', alignItems: 'center', gap: 1, borderBottom: '1px solid var(--color-divider)' }}>
              <IconButton size="small" onClick={() => handlePageChange(start - 1)} disabled={start <= 1}><ArrowBackIosNew fontSize="inherit" /></IconButton>
              <Box sx={{ fontSize: 12, color: 'var(--color-text-2)' }}>Section {start}–{end}</Box>
              <IconButton size="small" onClick={() => handlePageChange(end + 1)} disabled={end >= totalPages}><ArrowForwardIos fontSize="inherit" /></IconButton>
            </Box>
          )}
          <Box sx={{ flex: 1, overflow: 'auto', display: 'flex', flexDirection: 'column', alignItems: 'center', gap: 2, p: 2 }}>
            {pages.map((p) => {
              const idx = Math.max(0, Math.min((manifest.page_count || 1) - 1, p - 1));
              const pg = manifest.page_images[idx];
              if (!pg) return null;
              const url = `${API_BASE}/database/files/proxy?bucket=${encodeURIComponent(manifest.bucket || '')}&path=${encodeURIComponent(pg.full_image_path)}`;
              return (
                <ImageByProxy key={p} url={url} alt={`Page ${p}`} />
              );
            })}
          </Box>
        </Box>
      );
    }
    // Fallback legacy rendering
    if (!images.length) return <Box sx={{ p: 2 }}>No page images available.</Box>;
    return (
      <Box sx={{ width: '100%', height: '100%', overflow: 'auto', p: 2 }}>
        <Box sx={{ display: 'flex', flexDirection: 'column', alignItems: 'center', gap: 2 }}>
          {images.map((img, i) => (
            <img key={i} src={img.src} alt={`Page ${i + 1}`} style={{ boxShadow: '0 2px 8px rgba(0,0,0,0.15)', maxWidth: '100%' }} />
          ))}
        </Box>
      </Box>
    );
  }, [loading, error, images, manifest, pageObjectUrl, page, totalPages]);

  return (
    <Box sx={{ width: '100%', height: '100%', position: 'relative' }}>
      {content}
    </Box>
  );
};

export default CCDoclingViewer;

const ImageByProxy: React.FC<{ url: string; alt: string }> = ({ url, alt }) => {
  const { accessToken } = useAuth();
  const [blobUrl, setBlobUrl] = useState<string | undefined>(undefined);
  const [loading, setLoading] = useState<boolean>(true);
  const [error, setError] = useState<string | null>(null);
  useEffect(() => {
    let revoked: string | null = null;
    const load = async () => {
      try {
        const token = accessToken || '';
        const resp = await fetch(url, { headers: { Authorization: `Bearer ${token}` } });
        if (!resp.ok) throw new Error(`HTTP ${resp.status}`);
        const blob = await resp.blob();
        const obj = URL.createObjectURL(blob);
        setBlobUrl(obj);
        revoked = obj;
      } catch (e: unknown) {
        setError(e instanceof Error ? e.message : 'Failed to load page');
      } finally {
        setLoading(false);
      }
    };
    load();
    return () => { if (revoked) URL.revokeObjectURL(revoked); };
  }, [url]);

  if (loading) return <Box sx={{ p: 2 }}><CircularProgress size={18} /></Box>;
  if (error || !blobUrl) return <Box sx={{ p: 2, color: 'var(--color-text-2)' }}>{error || 'No image'}</Box>;
  return (
    <img src={blobUrl} alt={alt} style={{ maxWidth: '100%', height: 'auto', display: 'block', boxShadow: '0 2px 8px rgba(0,0,0,0.15)' }} />
  );
};