app/api/documents/route.js

import { HF_DATASET_BASE_URL, MAX_DOCS_TO_SCAN, getCorpus, getLinksRepoPath, getDocRepoPath, getDocLocalPath } from '../../../utils/config.js';
import yaml from 'js-yaml';

const isHFSpace = () => {
    return process.env.HF_TOKEN && process.env.NODE_ENV !== 'development';
};

export const dynamic = 'force-dynamic';


/**
 * Fetch annotator_config.yaml and return the doc list for a given user.
 * Returns null only if no username is provided (unauthenticated).
 * Returns empty {} if user is not in config (sees no docs).
 * Returns per-corpus assignments: { wbg: Set([1,2]), unhcr: Set([3,4]) }
 */
async function getUserAssignedDocs(username) {
    if (!username) return null;

    try {
        let text;
        if (isHFSpace()) {
            const configUrl = `${HF_DATASET_BASE_URL}/raw/main/annotation_data/annotator_config.yaml`;
            const res = await fetch(configUrl, {
                headers: { 'Authorization': `Bearer ${process.env.HF_TOKEN}` },
                cache: 'no-store'
            });
            if (!res.ok) return {}; // config missing — block access
            text = await res.text();
        } else {
            const fs = await import('fs');
            const path = await import('path');
            const filePath = path.join(process.cwd(), 'annotation_data', 'annotator_config.yaml');
            if (!fs.existsSync(filePath)) return {};
            text = fs.readFileSync(filePath, 'utf-8');
        }
        const config = yaml.load(text);

        const annotator = (config.annotators || []).find(a => a.username === username);
        if (!annotator || !annotator.docs) return null; // user not in config — see all docs

        // Support both old format (flat array) and new format (per-corpus object)
        if (Array.isArray(annotator.docs)) {
            // Legacy: flat array — treat as default corpus
            return { _flat: new Set(annotator.docs) };
        }

        // New format: { wbg: [1,2], unhcr: [3,4] }
        const result = {};
        for (const [corpusId, docList] of Object.entries(annotator.docs)) {
            if (Array.isArray(docList)) {
                result[corpusId] = new Set(docList);
            }
        }
        return result; // may be empty {} if user has no corpus assignments
    } catch (e) {
        console.warn('Could not load annotator_config.yaml:', e.message);
        return {}; // on error, block access rather than show all
    }
}

export async function GET(request) {
    try {
        const { searchParams } = new URL(request.url);
        const username = searchParams.get('user');

        const assignedDocs = null; // Bypassed assignment: show all successful documents to all users

        // Import corpora list
        const { getCorpora } = await import('../../../utils/config.js');
        const corpora = getCorpora();

        const allDocuments = [];

        for (const corpus of corpora) {
            // Determine which doc indices this user has for this corpus
            let userDocSet = null;
            if (assignedDocs) {
                if (assignedDocs._flat) {
                    // Legacy flat format — only applies to first/default corpus
                    userDocSet = corpus === corpora[0] ? assignedDocs._flat : new Set();
                } else {
                    userDocSet = assignedDocs[corpus.id] || new Set();
                }
                if (userDocSet.size === 0) continue; // no docs for this corpus
            }

            // Fetch the links file for this corpus
            let links;
            if (isHFSpace()) {
                const linksPath = getLinksRepoPath(corpus);
                const linksUrl = `${HF_DATASET_BASE_URL}/raw/main/${linksPath}`;
                const linksRes = await fetch(linksUrl, {
                    headers: { 'Authorization': `Bearer ${process.env.HF_TOKEN}` },
                    cache: 'no-store'
                });

                if (!linksRes.ok) {
                    console.warn(`No links file for corpus ${corpus.id}: ${linksRes.status}`);
                    continue;
                }
                links = await linksRes.json();
            } else {
                const fs = await import('fs');
                const path = await import('path');
                const filePath = path.join(process.cwd(), 'annotation_data', corpus.links_file);
                if (!fs.existsSync(filePath)) {
                    console.warn(`No links file for corpus ${corpus.id} locally: ${filePath}`);
                    continue;
                }
                links = JSON.parse(fs.readFileSync(filePath, 'utf-8'));
            }

            let successLinks = links
                .filter(l => l.status === 'success' && l.has_revalidation === true);

            if (userDocSet) {
                successLinks = successLinks.filter(l => userDocSet.has(l.index));
            }

            successLinks = successLinks.slice(0, MAX_DOCS_TO_SCAN);

            // Parallel fetch docs
            const results = await Promise.allSettled(
                successLinks.map(async (link) => {
                    let pagesData;
                    if (isHFSpace()) {
                        const docRepoPath = getDocRepoPath(corpus, link.index);
                        const docUrl = `${HF_DATASET_BASE_URL}/raw/main/${docRepoPath}`;
                        const docRes = await fetch(docUrl, {
                            headers: { 'Authorization': `Bearer ${process.env.HF_TOKEN}` },
                            cache: 'no-store'
                        });

                        if (!docRes.ok) return null;
                        pagesData = await docRes.json();
                    } else {
                        const fs = await import('fs');
                        const filePath = getDocLocalPath(corpus, link.index);
                        if (!fs.existsSync(filePath)) return null;
                        pagesData = JSON.parse(fs.readFileSync(filePath, 'utf-8'));
                    }
                    const allPages = pagesData
                        .map(page => page.document.pages[0]);
                    const pagesWithMentions = pagesData
                        .filter(page => {
                            if (!page.datasets) return false;
                            const validDatasets = page.datasets.filter(ds => {
                                if (ds.dataset_tag === 'non-dataset' && ds.dataset_name?.judge_agrees === true) {
                                    return false;
                                }
                                const confidence = ds.dataset_name?.confidence ?? 1.0;
                                if (confidence < 0.5) {
                                    return false;
                                }
                                return true;
                            });
                            return validDatasets.length > 0;
                        })
                        .map(page => page.document.pages[0]);

                    if (allPages.length === 0) return null;

                    const pdfUrl = link.direct_pdf_url;
                    if (!pdfUrl) return null;

                    return {
                        corpus: corpus.id,
                        corpus_name: corpus.name,
                        index: link.index,
                        original_name: link.original_name,
                        pdf_url: pdfUrl,
                        landing_page: link.landing_page_url,
                        annotatable_pages: allPages,
                        pages_with_mentions: pagesWithMentions
                    };
                })
            );

            const docs = results
                .filter(r => r.status === 'fulfilled' && r.value !== null)
                .map(r => r.value);

            allDocuments.push(...docs);
        }

        return new Response(JSON.stringify(allDocuments), {
            status: 200,
            headers: {
                'Content-Type': 'application/json',
                'Cache-Control': 'no-store, max-age=0, must-revalidate'
            }
        });
    } catch (error) {
        console.error(error);
        return new Response(
            JSON.stringify({ error: "Failed to fetch documents" }),
            { status: 500, headers: { 'Content-Type': 'application/json' } }
        );
    }
}