import os, json, gc, pickle, traceback, zipfile, shutil, uuid, logging
from datetime import datetime
from flask import Flask, render_template, request, jsonify, send_file
from werkzeug.utils import secure_filename
import numpy as np
import faiss

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = Flask(__name__)
app.config['MAX_CONTENT_LENGTH'] = 500 * 1024 * 1024   # 500 MB
app.config['TEMP_FOLDER'] = os.path.join(os.path.expanduser('~'), 'tmp_merge')
os.makedirs(app.config['TEMP_FOLDER'], exist_ok=True)

@app.route('/')
def merge_page():
    return render_template('merge.html')

@app.route('/api/merge', methods=['POST'])
def merge_zips():
    zip_files = request.files.getlist('files')
    out_path = request.form.get('out_path', '/tmp').strip()
    out_name = request.form.get('out_name', 'merged_vectorstore').strip()
    if out_name.lower().endswith('.zip'):
        out_name = out_name[:-4]

    if not zip_files:
        return jsonify({'error': 'No files uploaded'}), 400
    for f in zip_files:
        if not f.filename.lower().endswith('.zip'):
            return jsonify({'error': f'{f.filename} is not a .zip file'}), 400

    os.makedirs(out_path, exist_ok=True)

    run_id = uuid.uuid4().hex[:8]
    work_dir = os.path.join(app.config['TEMP_FOLDER'], f'merge_{run_id}')
    os.makedirs(work_dir, exist_ok=True)

    main_index = None
    all_documents = []
    all_metadata = []
    embedding_dim = None
    total_merged = 0
    errors, skipped = [], []

    try:
        for idx, zip_file in enumerate(zip_files, 1):
            safe_name = secure_filename(zip_file.filename)
            zip_save = os.path.join(work_dir, safe_name)
            extract_dir = os.path.join(work_dir, f'ex_{idx:02d}')
            os.makedirs(extract_dir, exist_ok=True)

            zip_file.save(zip_save)
            if not zipfile.is_zipfile(zip_save):
                errors.append(f'{safe_name}: not a valid zip — skipped')
                continue

            with zipfile.ZipFile(zip_save, 'r') as zf:
                zf.extractall(extract_dir)

            # Locate the three required files
            index_file = data_file = meta_file = None
            for root, dirs, fnames in os.walk(extract_dir):
                for fn in fnames:
                    full = os.path.join(root, fn)
                    if fn == 'faiss.index' and index_file is None: index_file = full
                    if fn == 'data.pkl'    and data_file  is None: data_file  = full
                    if fn == 'meta.json'   and meta_file  is None: meta_file  = full

            if not index_file or not data_file:
                errors.append(f'{safe_name}: missing faiss.index/data.pkl — skipped')
                continue

            # Load meta.json (optional, we can create it later)
            meta = {}
            if meta_file:
                with open(meta_file, 'r') as f:
                    meta = json.load(f)

            try:
                sub_index = faiss.read_index(index_file)
                with open(data_file, 'rb') as f:
                    sub_data = pickle.load(f)
            except Exception as e:
                errors.append(f'{safe_name}: load error — {e}')
                continue

            sub_docs = sub_data.get('documents', [])
            sub_meta = sub_data.get('metadata', [])
            n = sub_index.ntotal

            if n == 0:
                skipped.append(f'{safe_name}: empty index — skipped')
                continue

            if n != len(sub_docs) or n != len(sub_meta):
                errors.append(f'{safe_name}: index/pkl mismatch ({n} vectors, {len(sub_docs)} docs, {len(sub_meta)} meta) — skipped')
                continue

            # Dimension check
            dim = sub_index.d
            if embedding_dim is None:
                embedding_dim = dim
                main_index = faiss.IndexFlatL2(dim)
            elif dim != embedding_dim:
                errors.append(f'{safe_name}: dimension mismatch (got {dim}, expected {embedding_dim}) — skipped')
                continue

            # Extract vectors
            vectors = np.zeros((n, dim), dtype=np.float32)
            sub_index.reconstruct_n(0, n, vectors)
            main_index.add(vectors)
            all_documents.extend(sub_docs)
            all_metadata.extend(sub_meta)
            total_merged += n

            logger.info(f'  ✓ {safe_name}: +{n} vectors (total {total_merged})')
            del sub_index, vectors, sub_docs, sub_meta, sub_data
            gc.collect()

        if total_merged == 0:
            return jsonify({'error': 'No vectors merged', 'errors': errors, 'skipped': skipped}), 400

        # Prepare output zip
        merged_dir = os.path.join(work_dir, 'merged')
        os.makedirs(merged_dir, exist_ok=True)

        # Write faiss.index
        index_out = os.path.join(merged_dir, 'faiss.index')
        faiss.write_index(main_index, index_out)

        # Write data.pkl
        data_out = os.path.join(merged_dir, 'data.pkl')
        with open(data_out, 'wb') as f:
            pickle.dump({'documents': all_documents, 'metadata': all_metadata}, f)

        # Write fresh meta.json
        new_meta = {
            'version': 2,
            'similarity': 'cosine',   # safe default; adjust if you know the real one
            'embedding_dim': embedding_dim,
            'chunk_count': total_merged
        }
        meta_out = os.path.join(merged_dir, 'meta.json')
        with open(meta_out, 'w') as f:
            json.dump(new_meta, f)

        # Package into zip
        final_zip = os.path.join(out_path, out_name + '.zip')
        with zipfile.ZipFile(final_zip, 'w', zipfile.ZIP_DEFLATED) as zf:
            zf.write(index_out, 'faiss.index')
            zf.write(data_out,  'data.pkl')
            zf.write(meta_out,  'meta.json')

        return send_file(final_zip, as_attachment=True, download_name=os.path.basename(final_zip))

    except Exception as e:
        logger.error(f'Merge error: {e}')
        logger.error(traceback.format_exc())
        return jsonify({'error': str(e)}), 500

    finally:
        if os.path.exists(work_dir):
            shutil.rmtree(work_dir, ignore_errors=True)


if __name__ == '__main__':
    port = int(os.environ.get('PORT', 7860))
    logger.info(f'✅ Starting merge server on port {port}')
    app.run(host='0.0.0.0', port=port, debug=False)