import os
import uuid
import fitz
import base64
from datetime import datetime
from flask import current_app, request, jsonify, render_template, redirect, url_for, send_file
from werkzeug.utils import secure_filename
from .common import main_bp, get_db_connection, login_required, current_user
from utils import create_a4_pdf_from_images
from processing import create_pdf_from_full_images
from resize import expand_pdf_for_notes
from database import get_folder_tree, get_all_descendant_folder_ids
from strings import ROUTE_GENERATE_PDF, METHOD_POST

@main_bp.route('/generate_color_rm_pdf/<session_id>')
@login_required
def generate_color_rm_pdf(session_id):
    conn = get_db_connection()
    session_data = conn.execute('SELECT user_id, original_filename FROM sessions WHERE id = ?', (session_id,)).fetchone()
    if not session_data or session_data['user_id'] != current_user.id:
        conn.close(); return jsonify({'error': 'Unauthorized'}), 403
    start_page, end_page = request.args.get('start', type=int), request.args.get('end', type=int)
    query, params = "SELECT * FROM images WHERE session_id = ? AND image_type = 'original'", [session_id]
    if start_page: query += " AND image_index >= ?"; params.append(start_page - 1)
    if end_page: query += " AND image_index <= ?"; params.append(end_page - 1)
    query += " ORDER BY image_index"
    images = conn.execute(query, params).fetchall()
    pdf_image_paths = []
    for img in images:
        processed = conn.execute("SELECT processed_filename FROM images WHERE session_id = ? AND image_index = ? AND image_type = 'color_rm'", (session_id, img['image_index'])).fetchone()
        path = os.path.join(current_app.config['PROCESSED_FOLDER'], processed['processed_filename']) if processed and processed['processed_filename'] else os.path.join(current_app.config['UPLOAD_FOLDER'], img['filename'])
        if os.path.exists(path): pdf_image_paths.append(path)
    if not pdf_image_paths: conn.close(); return "No images found", 404
    range_suffix = f"_pg{start_page or 1}-{end_page or 'end'}" if start_page or end_page else ""
    pdf_filename = f"Color_Removed_{session_data['original_filename']}{range_suffix}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.pdf"
    output_path = os.path.join(current_app.config['OUTPUT_FOLDER'], pdf_filename)
    res = getattr(current_user, 'color_rm_dpi', 200.0)
    if create_pdf_from_full_images(pdf_image_paths, output_path, resolution=res):
        conn.execute('INSERT INTO generated_pdfs (session_id, filename, subject, user_id) VALUES (?, ?, ?, ?)', (session_id, pdf_filename, "Color Removal Export", current_user.id))
        conn.commit(); conn.close(); return redirect(url_for('main.pdf_manager'))
    else: conn.close(); return "Failed", 500

@main_bp.route('/generate_preview', methods=[METHOD_POST])
@login_required
def generate_preview():
    data = request.json
    sid = data['session_id']
    conn = get_db_connection()
    session_owner = conn.execute('SELECT user_id FROM sessions WHERE id = ?', (sid,)).fetchone()
    if not session_owner or session_owner['user_id'] != current_user.id:
        conn.close(); return jsonify({'error': 'Unauthorized'}), 403
    all_questions = [dict(row) for row in conn.execute("""
        SELECT q.*, i.filename, i.processed_filename, i.note_filename, i.include_note_in_pdf
        FROM questions q JOIN images i ON q.image_id = i.id
        WHERE q.session_id = ? ORDER BY i.id
    """, (sid,)).fetchall()]
    conn.close()
    # Only include note_filename if include_note_in_pdf is true
    for q in all_questions:
        if not q.get('include_note_in_pdf', 1):
            q['note_filename'] = None
    all_questions.extend(data.get('miscellaneous_questions', []))
    filter_type = data.get('filter_type', 'all')
    filtered = [q for q in all_questions if filter_type == 'all' or q['status'] == filter_type]
    if not filtered: return jsonify({'error': 'No questions match'}), 400
    images_per_page = int(data.get('images_per_page', 4))
    practice_mode = data.get('practice_mode', 'none')
    practice_mode_settings = {'portrait_2': {'images_per_page': 2, 'orientation': 'portrait', 'grid_rows': 2, 'grid_cols': 1}, 'portrait_3': {'images_per_page': 3, 'orientation': 'portrait', 'grid_rows': 3, 'grid_cols': 1}, 'landscape_2': {'images_per_page': 2, 'orientation': 'landscape', 'grid_rows': 2, 'grid_cols': 1}, 'portrait_2_spacious': {'images_per_page': 2, 'orientation': 'portrait', 'grid_rows': 2, 'grid_cols': 1}}
    if practice_mode in practice_mode_settings:
        s = practice_mode_settings[practice_mode]
        ipp, orientation, rows, cols = s['images_per_page'], s['orientation'], s['grid_rows'], s['grid_cols']
    else:
        ipp, orientation = int(data.get('images_per_page', 4)), data.get('orientation', 'portrait')
        rows = int(data['grid_rows']) if data.get('grid_rows') else None
        cols = int(data['grid_cols']) if data.get('grid_cols') else None
    pdf_bytes = create_a4_pdf_from_images(filtered[:ipp], current_app.config['PROCESSED_FOLDER'], None, ipp, output_folder=None, orientation=orientation, grid_rows=rows, grid_cols=cols, practice_mode=practice_mode, return_bytes=True, font_size_scale=float(data.get('font_size_scale', 1.0)))
    if pdf_bytes:
        try:
            doc = fitz.open(stream=pdf_bytes, filetype="pdf")
            pix = doc.load_page(0).get_pixmap(dpi=150)
            img_b64 = base64.b64encode(pix.tobytes("png")).decode('utf-8')
            return jsonify({'success': True, 'preview_image': f'data:image/png;base64,{img_b64}'})
        except Exception as e: return jsonify({'error': str(e)}), 500
        finally: 
            if 'doc' in locals(): doc.close()
    return jsonify({'error': 'Failed'}), 500

@main_bp.route(ROUTE_GENERATE_PDF, methods=[METHOD_POST])
@login_required
def generate_pdf():
    data = request.json
    sid = data['session_id']
    conn = get_db_connection()
    session_owner = conn.execute('SELECT user_id FROM sessions WHERE id = ?', (sid,)).fetchone()
    if not session_owner or session_owner['user_id'] != current_user.id:
        conn.close(); return jsonify({'error': 'Unauthorized'}), 403
    all_questions = [dict(row) for row in conn.execute("""
        SELECT q.*, i.filename, i.processed_filename, i.note_filename, i.include_note_in_pdf
        FROM questions q JOIN images i ON q.image_id = i.id
        WHERE q.session_id = ? ORDER BY i.id
    """, (sid,)).fetchall()]
    # Only include note_filename if include_note_in_pdf is true
    for q in all_questions:
        if not q.get('include_note_in_pdf', 1):
            q['note_filename'] = None
    all_questions.extend(data.get('miscellaneous_questions', []))
    filter_type = data.get('filter_type', 'all')
    filtered = [q for q in all_questions if filter_type == 'all' or q['status'] == filter_type]
    if not filtered: conn.close(); return jsonify({'error': 'No questions match'}), 400
    pdf_filename = f"{secure_filename(data.get('pdf_name', 'analysis'))}_{datetime.now().strftime('%Y%m%d_%H%M')}.pdf"
    practice_mode = data.get('practice_mode', 'none')
    practice_mode_settings = {'portrait_2': {'images_per_page': 2, 'orientation': 'portrait', 'grid_rows': 2, 'grid_cols': 1}, 'portrait_3': {'images_per_page': 3, 'orientation': 'portrait', 'grid_rows': 3, 'grid_cols': 1}, 'landscape_2': {'images_per_page': 2, 'orientation': 'landscape', 'grid_rows': 2, 'grid_cols': 1}, 'portrait_2_spacious': {'images_per_page': 2, 'orientation': 'portrait', 'grid_rows': 2, 'grid_cols': 1}}
    if practice_mode in practice_mode_settings:
        s = practice_mode_settings[practice_mode]; ipp, orientation, rows, cols = s['images_per_page'], s['orientation'], s['grid_rows'], s['grid_cols']
    else:
        ipp, orientation = int(data.get('images_per_page', 4)), data.get('orientation', 'portrait')
        rows = int(data['grid_rows']) if data.get('grid_rows') else None
        cols = int(data['grid_cols']) if data.get('grid_cols') else None
    if create_a4_pdf_from_images(filtered, current_app.config['PROCESSED_FOLDER'], pdf_filename, ipp, current_app.config['OUTPUT_FOLDER'], orientation, rows, cols, practice_mode, font_size_scale=float(data.get('font_size_scale', 1.0))):
        s_info = conn.execute('SELECT original_filename FROM sessions WHERE id = ?', (sid,)).fetchone()
        conn.execute('INSERT INTO generated_pdfs (session_id, filename, subject, tags, notes, source_filename, user_id) VALUES (?, ?, ?, ?, ?, ?, ?)', (sid, pdf_filename, data.get('subject'), data.get('tags'), data.get('notes'), s_info['original_filename'] if s_info else 'Unknown', current_user.id))
        conn.commit(); conn.close(); return jsonify({'success': True, 'pdf_filename': pdf_filename})
    else: conn.close(); return jsonify({'error': 'Failed'}), 500

@main_bp.route('/merge_pdfs', methods=[METHOD_POST])
@login_required
def merge_pdfs():
    data = request.json
    pdf_ids = data.get('pdf_ids', [])
    if len(pdf_ids) < 2: return jsonify({'error': 'Select at least two.'}), 400
    try:
        conn = get_db_connection()
        placeholders = ', '.join('?' * len(pdf_ids))
        pdfs = conn.execute(f"SELECT filename FROM generated_pdfs WHERE id IN ({placeholders}) AND user_id = ?", (*pdf_ids, current_user.id)).fetchall()
        if len(pdfs) != len(pdf_ids): conn.close(); return jsonify({'error': 'Not found/Unauthorized'}), 404
        merged = fitz.open()
        src_fnames = []
        for p in pdfs:
            src_fnames.append(p['filename'])
            path = os.path.join(current_app.config['OUTPUT_FOLDER'], p['filename'])
            if os.path.exists(path):
                with fitz.open(path) as d: merged.insert_pdf(d)
        new_name = f"merged_{datetime.now().strftime('%Y%m%d_%H%M%S')}.pdf"
        merged.save(os.path.join(current_app.config['OUTPUT_FOLDER'], new_name)); merged.close()
        sid = str(uuid.uuid4())
        conn.execute('INSERT INTO sessions (id, original_filename, user_id) VALUES (?, ?, ?)', (sid, f"Merged from {len(src_fnames)} files", current_user.id))
        conn.execute('INSERT INTO generated_pdfs (session_id, filename, subject, tags, notes, source_filename, user_id) VALUES (?, ?, ?, ?, ?, ?, ?)', (sid, new_name, "Merged Document", 'merged', f"Merged from:\n" + "\n".join(src_fnames), ", ".join(src_fnames), current_user.id))
        conn.commit(); conn.close(); return jsonify({'success': True, 'new_filename': new_name})
    except Exception as e: return jsonify({'error': str(e)}), 500

@main_bp.route('/resize/', defaults={'folder_path': ''}, methods=['GET', 'POST'])
@main_bp.route('/resize/browse/<path:folder_path>', methods=['GET', 'POST'])
@login_required
def resize_pdf_route(folder_path):
    if request.method == 'POST':
        in_name, out_name = request.form.get('input_pdf'), request.form.get('output_pdf')
        bg, pattern, patt_col = request.form.get('bg_color', '#FFFFFF'), request.form.get('pattern'), request.form.get('pattern_color', '#CCCCCC')
        mode, direction, add_space = request.form.get('mode', 'notes_only'), request.form.get('stitch_direction', 'horizontal'), 'add_space' in request.form
        conn = get_db_connection()
        if not conn.execute('SELECT user_id FROM generated_pdfs WHERE filename = ? AND user_id = ?', (in_name, current_user.id)).fetchone():
            conn.close(); return "Unauthorized", 403
        def htr(h): h = h.lstrip('#'); return tuple(int(h[i:i+2], 16) / 255.0 for i in (0, 2, 4))
        try:
            expand_pdf_for_notes(os.path.join(current_app.config['OUTPUT_FOLDER'], in_name), os.path.join(current_app.config['OUTPUT_FOLDER'], out_name), bg_color=htr(bg), mode=mode, stitch_direction=direction, add_space=add_space, pattern=pattern, pattern_color=htr(patt_col))
            sid = str(uuid.uuid4())
            conn.execute('INSERT INTO sessions (id, original_filename, user_id) VALUES (?, ?, ?)', (sid, f"Resized from {in_name}", current_user.id))
            conn.execute('INSERT INTO generated_pdfs (session_id, filename, subject, tags, notes, source_filename, user_id) VALUES (?, ?, ?, ?, ?, ?, ?)', (sid, out_name, f"Resized - {os.path.basename(in_name)}", 'resized', f"mode={mode}, direction={direction}, space={add_space}, bg={bg}, pattern={pattern}", in_name, current_user.id))
            conn.commit(); conn.close(); return redirect(url_for('main.pdf_manager'))
        except Exception as e: conn.close(); return str(e), 500
    else:
        conn = get_db_connection()
        search, recursive = request.args.get('search', ''), request.args.get('recursive') == 'true'
        q_params, where = [current_user.id], ['user_id = ?']
        if search: where.append('(filename LIKE ? OR subject LIKE ? OR tags LIKE ?)'); st = f'%{search}%'; q_params.extend([st, st, st])
        f_id, sub, crumbs = None, [], []
        if folder_path:
            p_id = None
            for i, part in enumerate(folder_path.split('/')):
                res = conn.execute("SELECT id FROM folders WHERE name = ? AND user_id = ? AND (parent_id = ? OR (? IS NULL AND parent_id IS NULL))", (part, current_user.id, p_id, p_id)).fetchone()
                if not res: return redirect(url_for('main.resize_pdf_route'))
                p_id = res['id']; crumbs.append({'name': part, 'path': '/'.join(folder_path.split('/')[:i+1])})
            f_id = p_id
        if recursive and search:
            if f_id:
                all_f_ids = [f_id] + get_all_descendant_folder_ids(conn, f_id, current_user.id)
                where.append(f"folder_id IN ({', '.join('?' * len(all_f_ids))})"); q_params.extend(all_f_ids)
        else:
            if f_id: where.append('folder_id = ?'); q_params.append(f_id)
            else: where.append('folder_id IS NULL')
        if f_id: sub = conn.execute('SELECT * FROM folders WHERE parent_id = ? AND user_id = ? ORDER BY name', (f_id, current_user.id)).fetchall()
        else: sub = conn.execute('SELECT * FROM folders WHERE parent_id IS NULL AND user_id = ? ORDER BY name', (current_user.id,)).fetchall()
        base_q = 'SELECT * FROM generated_pdfs WHERE ' + ' AND '.join(where) + ' ORDER BY created_at DESC'
        pdfs = conn.execute(base_q, q_params).fetchall()
        tree = get_folder_tree(user_id=current_user.id); conn.close()
        return render_template('resize.html', pdfs=[dict(row) for row in pdfs], subfolders=[dict(row) for row in sub], current_folder_id=f_id, breadcrumbs=crumbs, folder_tree=tree, search_query=search, recursive=recursive)

@main_bp.route('/print_pdfs', methods=['POST'])
@login_required
def print_pdfs():
    pdf_ids = request.form.getlist('pdf_ids')
    if not pdf_ids: return jsonify({'error': 'None selected'}), 400
    conn = get_db_connection()
    placeholders = ','.join('?' for _ in pdf_ids)
    pdfs = conn.execute(f"SELECT filename, subject FROM generated_pdfs WHERE id IN ({placeholders}) AND user_id = ?", (*pdf_ids, current_user.id)).fetchall()
    conn.close()
    if not pdfs: return jsonify({'error': 'Not found'}), 404
    merged = fitz.open()
    for p in pdfs:
        path = os.path.join(current_app.config['OUTPUT_FOLDER'], p['filename'])
        if os.path.exists(path):
            with fitz.open(path) as d:
                if len(d) > 0:
                    try:
                        # Use insert_text which is simpler and less prone to font errors
                        d[0].insert_text(fitz.Point(20, 15), f"Subject: {p['subject']} | File: {p['filename']}", fontsize=6, color=(0, 0, 0))
                    except Exception as e:
                        current_app.logger.error(f"Could not insert header for {p['filename']}: {e}")
                merged.insert_pdf(d)
    tmp_name = f'printed_documents_{datetime.now().strftime("%Y%m%d%H%M%S")}.pdf'
    tmp_path = os.path.join(current_app.config['TEMP_FOLDER'], tmp_name)
    os.makedirs(current_app.config['TEMP_FOLDER'], exist_ok=True); merged.save(tmp_path); merged.close()
    return jsonify({'success': True, 'url': url_for('main.view_generated_pdf', filename=tmp_name)})