import os
import random
import io
from flask import Flask, render_template, request, jsonify, send_file
from reportlab.pdfgen import canvas
from reportlab.lib.pagesizes import A4
from reportlab.pdfbase import pdfmetrics
from reportlab.pdfbase.ttfonts import TTFont
from reportlab.lib.units import cm

app = Flask(__name__)

# --- Font Registration ---
FONT_NAME = "Helvetica" # Default fallback
CN_FONT_NAME = None

def register_fonts():
    global FONT_NAME, CN_FONT_NAME
    # Priority list for Chinese fonts
    font_paths = [
        "/usr/share/fonts/truetype/wqy/wqy-microhei.ttc",  # Debian/Ubuntu (Docker)
        "/usr/share/fonts/truetype/wqy/wqy-zenhei.ttc",
        "./static/fonts/SimHei.ttf",                       # Local custom
        "/System/Library/Fonts/STHeiti Light.ttc",         # macOS fallback (optional)
    ]
    
    for path in font_paths:
        if os.path.exists(path):
            try:
                # Try to register. For TTC, usually need to try index 0
                # If it fails, reportlab might raise error
                # Note: 'wqy-microhei' is usually "WenQuanYi Micro Hei"
                font_key = "ChineseFont"
                pdfmetrics.registerFont(TTFont(font_key, path))
                CN_FONT_NAME = font_key
                print(f"Successfully registered Chinese font: {path}")
                break
            except Exception as e:
                print(f"Failed to register font {path}: {e}")

register_fonts()

# --- Math Logic ---

def generate_problem(ops, min_val, max_val, allow_remainder=False, no_negative=True, blank_prob=0):
    """
    Generates a single math problem.
    blank_prob: 0.0 to 1.0, probability of having a blank in operands (e.g. 3 + _ = 5)
    """
    op = random.choice(ops)
    question = ""
    answer = ""
    
    # 1. Generate Numbers based on Operation
    if op == '+':
        a = random.randint(min_val, max_val)
        b = random.randint(min_val, max_val)
        res = a + b
        
        # Format: a + b = res
        # With blank_prob, we might hide a or b
        if random.random() < blank_prob:
            # Hide a or b
            if random.choice([True, False]):
                question = f"(   ) + {b} = {res}"
                answer = str(a)
            else:
                question = f"{a} + (   ) = {res}"
                answer = str(b)
        else:
            question = f"{a} + {b} ="
            answer = str(res)
            
    elif op == '-':
        a = random.randint(min_val, max_val)
        b = random.randint(min_val, max_val)
        if no_negative and a < b:
            a, b = b, a
        res = a - b
        
        if random.random() < blank_prob:
            if random.choice([True, False]):
                question = f"(   ) - {b} = {res}"
                answer = str(a)
            else:
                question = f"{a} - (   ) = {res}"
                answer = str(b)
        else:
            question = f"{a} - {b} ="
            answer = str(res)
            
    elif op == '×':
        a = random.randint(min_val, max_val)
        b = random.randint(min_val, max_val)
        res = a * b
        
        if random.random() < blank_prob:
            if random.choice([True, False]):
                question = f"(   ) × {b} = {res}"
                answer = str(a)
            else:
                question = f"{a} × (   ) = {res}"
                answer = str(b)
        else:
            question = f"{a} × {b} ="
            answer = str(res)
            
    elif op == '÷':
        b = random.randint(min_val, max_val)
        if b == 0: b = 1
        
        if allow_remainder:
            # Division with remainder allowed
            # We construct 'a' such that it might have remainder
            # But usually we want controlled difficulty.
            # Let's pick 'res' (quotient) and 'rem' (remainder)
            quotient = random.randint(min_val, max_val)
            remainder = random.randint(0, b - 1)
            a = quotient * b + remainder
            res = f"{quotient} ... {remainder}" if remainder != 0 else str(quotient)
            
            # Blanks in division with remainder is complex, let's skip blanks for remainder mode for simplicity
            # or only hide dividend/divisor if remainder is 0
            question = f"{a} ÷ {b} ="
            answer = str(res)
            
        else:
            # Exact division
            res_val = random.randint(min_val, max_val)
            a = res_val * b
            res = str(res_val)
            
            if random.random() < blank_prob:
                if random.choice([True, False]):
                    question = f"(   ) ÷ {b} = {res}"
                    answer = str(a)
                else:
                    # a / ? = res => ? = a / res
                    # Avoid division by zero if res is 0
                    if res_val == 0:
                         question = f"{a} ÷ {b} =" # Fallback
                         answer = str(res)
                    else:
                        question = f"{a} ÷ (   ) = {res}"
                        answer = str(b)
            else:
                question = f"{a} ÷ {b} ="
                answer = str(res)
                
    return {"question": question, "answer": answer}

def generate_sheet_data(count, ops, min_val, max_val, blank_prob=0):
    problems = []
    for _ in range(count):
        p = generate_problem(ops, min_val, max_val, blank_prob=blank_prob)
        problems.append(p)
    return problems

# --- PDF Logic ---

def draw_header(c, width, height, margin_y, title):
    # Use Chinese font if available, else Helvetica
    title_font = CN_FONT_NAME if CN_FONT_NAME else "Helvetica-Bold"
    text_font = CN_FONT_NAME if CN_FONT_NAME else "Helvetica"
    
    # Title
    c.setFont(title_font, 24)
    # If no Chinese font and title contains non-ascii, fallback to English
    display_title = title
    if not CN_FONT_NAME and any(ord(char) > 127 for char in title):
        display_title = "Math Worksheet"
        
    c.drawCentredString(width/2, height - margin_y, display_title)
    
    # Info Line
    c.setFont(text_font, 12)
    info_y = height - margin_y - 1.2*cm
    
    if CN_FONT_NAME:
        c.drawString(2*cm, info_y, "姓名: ______________")
        c.drawString(width/2 - 2*cm, info_y, "日期: ______________")
        c.drawString(width - 5*cm, info_y, "用时: ________")
        c.drawString(width - 2.5*cm, info_y, "分数: ________")
    else:
        c.drawString(2*cm, info_y, "Name: ______________   Date: ______________   Time: ______   Score: ______")

def create_pdf(problems, title="口算练习题", with_answers=True, cols=4):
    buffer = io.BytesIO()
    c = canvas.Canvas(buffer, pagesize=A4)
    width, height = A4
    
    # Config
    margin_x = 2 * cm
    margin_y = 2 * cm
    
    # Calculate layout
    col_width = (width - 2 * margin_x) / cols
    row_height = 1.4 * cm
    
    # --- Problem Pages ---
    
    # Draw first header
    draw_header(c, width, height, margin_y, title)
    
    y_start = height - margin_y - 2.5*cm
    y = y_start
    
    # Font for problems - Use Helvetica for numbers/symbols as it looks better usually
    # But if question contains Chinese (not now), we need CN font.
    # Our questions are "1 + 2 =", pure ascii usually.
    # But we might have brackets ( ) which are fine.
    problem_font = "Helvetica" 
    c.setFont(problem_font, 14)
    
    current_col = 0
    
    # We might need multiple pages for problems
    for i, p in enumerate(problems):
        x = margin_x + current_col * col_width
        
        c.drawString(x, y, p['question'])
        
        current_col += 1
        if current_col >= cols:
            current_col = 0
            y -= row_height
            
        # Check page break
        if y < margin_y:
            c.showPage()
            # New page header (Simplified)
            y = height - margin_y
            c.setFont(problem_font, 14)
            
    # End of problems
    c.showPage()
    
    # --- Answer Key Page (Optional) ---
    if with_answers:
        # Header
        title_font = CN_FONT_NAME if CN_FONT_NAME else "Helvetica-Bold"
        c.setFont(title_font, 18)
        ans_title = "参考答案" if CN_FONT_NAME else "Answer Key"
        c.drawCentredString(width/2, height - margin_y, ans_title)
        
        y = height - margin_y - 1.5*cm
        c.setFont("Helvetica", 11) # Smaller font for answers
        
        # Use more cols for answers to save paper
        ans_cols = 5
        ans_col_width = (width - 2 * margin_x) / ans_cols
        ans_row_height = 0.8 * cm
        
        current_col = 0
        
        for i, p in enumerate(problems):
            x = margin_x + current_col * ans_col_width
            
            # Format: "1) 15"
            ans_text = f"{i+1}) {p['answer']}"
            c.drawString(x, y, ans_text)
            
            current_col += 1
            if current_col >= ans_cols:
                current_col = 0
                y -= ans_row_height
                
            if y < margin_y:
                c.showPage()
                y = height - margin_y
                c.setFont("Helvetica", 11)
        
        c.showPage()

    c.save()
    buffer.seek(0)
    return buffer

# --- Routes ---

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

@app.route('/api/preview', methods=['POST'])
def preview():
    data = request.json
    count = int(data.get('count', 20))
    min_val = int(data.get('min', 1))
    max_val = int(data.get('max', 20))
    ops = data.get('ops', ['+'])
    blank_prob = float(data.get('blank_prob', 0)) # 0 to 1
    
    problems = generate_sheet_data(count, ops, min_val, max_val, blank_prob)
    return jsonify(problems)

@app.route('/api/download', methods=['POST'])
def download():
    data = request.json
    count = int(data.get('count', 50))
    min_val = int(data.get('min', 1))
    max_val = int(data.get('max', 20))
    ops = data.get('ops', ['+'])
    blank_prob = float(data.get('blank_prob', 0))
    with_answers = data.get('with_answers', False)
    cols = int(data.get('cols', 4))
    
    problems = generate_sheet_data(count, ops, min_val, max_val, blank_prob)
    pdf_buffer = create_pdf(problems, with_answers=with_answers, cols=cols)
    
    return send_file(
        pdf_buffer,
        as_attachment=True,
        download_name='math_worksheet.pdf',
        mimetype='application/pdf'
    )

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7860)