"""app.py - Quart (async Flask) version with native WebSocket + asyncssh task runner.
Same API as before, but fully async for speed parity with original panel."""
import os, json, time, random, string, secrets, re, asyncio, base64
from datetime import datetime
from quart import Quart, request, jsonify, session, send_file, render_template, send_from_directory, redirect, websocket
from cards import register_user, renew_user, login_user, logout_user, current_user, countdown_from_expires, is_expired
from store import get_options, save_options, get_servers, save_servers, clear_servers, add_servers, update_server, get_templates, save_templates, add_template, delete_template, set_active_template, get_files_meta, save_files_meta, gen_card_code, add_card, get_cards
from dkim_modes import DKIM_MODES, resolve_dkim
from task_runner import runner, TASK_LABELS

app = Quart(__name__, template_folder="templates", static_folder="static")
app.secret_key = "mailpanel_secret_change_me_2026"

def get_full_options():
    from options_meta import OPTIONS_META
    opts = get_options()
    merged = {}
    merged.update(OPTIONS_META)
    merged.update(opts)
    for k in OPTIONS_META:
        if isinstance(OPTIONS_META[k], list) and OPTIONS_META[k] and isinstance(OPTIONS_META[k][0], dict):
            merged[k] = OPTIONS_META[k]
        elif k in ("dkim_custom_headers_default","strip_headers_default","received_ehlo_builtin_count",
                    "has_saved_servers","supported_os_label"):
            merged[k] = OPTIONS_META[k]
    return merged

def add_computed_labels(opts):
    from options_meta import OPTIONS_META
    for s in opts.get("mail_stacks", []):
        if s["id"] == opts.get("mail_stack", ""):
            opts["mail_stack_label"] = s.get("label", "")
    for s in opts.get("acme_cas", []):
        if s["id"] == opts.get("acme_ca", ""):
            opts["acme_ca_label"] = s.get("label", "")
    for s in opts.get("dkim_modes", []):
        if s["id"] == opts.get("dkim_sign_mode", ""):
            opts["dkim_sign_mode_label"] = s.get("label", "")
    opts["allowed_mail_stacks"] = [s["id"] for s in opts.get("mail_stacks", [])]
    opts["monitor_max_age_days"] = 0.0
    opts["monitor_log_lines"] = 0
    cur_stack = opts.get("mail_stack", "haraka")
    stack_dkim_first = {
        "postfix": "postfix_opendkim_official", "postfix_outbound": "postfix_opendkim_official",
        "haraka": "haraka_dkim_official", "maddy": "maddy_extended",
        "opensmtpd": "haraka_dkim_official", "exim4": "exim4_extended",
        "exim4_outbound": "exim4_extended", "courier_mta": "haraka_dkim_official",
        "poste_io": "haraka_dkim_official", "stalwart": "stalwart_extended",
    }
    first_mode_id = stack_dkim_first.get(cur_stack, "haraka_dkim_official")
    dkim_modes = opts.get("dkim_modes", [])
    new_order = [m for m in dkim_modes if m["id"] == first_mode_id]
    new_order += [m for m in dkim_modes if m["id"] != first_mode_id]
    opts["dkim_modes"] = new_order
    hsc = opts.get("header_strip_choices", [])
    if cur_stack != "haraka":
        hsc = [h for h in hsc if h.get("source") != "stack"]
    opts["header_strip_choices"] = hsc
    shd = opts.get("strip_headers_default", "")
    stack_extra = ""
    if cur_stack == "haraka":
        stack_extra = ", X-深情, X-Haraka-Latency, X-Haraka-Auth-User"
    if stack_extra and stack_extra.strip(", ") not in shd:
        opts["strip_headers_default"] = shd + stack_extra
    profiles = opts.get("header_strip_smart_profiles", [])
    stack_headers_to_add = []
    if cur_stack == "haraka":
        stack_headers_to_add = ["X-深情", "X-Haraka-Latency", "X-Haraka-Auth-User"]
    for prof in profiles:
        headers = list(prof.get("headers", []))
        for h in stack_headers_to_add:
            if h not in headers:
                headers.append(h)
        prof["headers"] = headers
    if "quick_paste" in opts and isinstance(opts["quick_paste"], dict):
        opts["quick_paste"].pop("password", None)
    return opts

def get_full_options_with_labels():
    opts = get_full_options()
    return add_computed_labels(opts)

# ===== REST API =====
@app.route("/api/auth/login", methods=["POST"])
async def api_login():
    data = await request.get_json(force=True)
    from cards import hash_password, verify_password, is_expired, countdown_from_expires, gen_expires_at
    from store import get_user
    username = data.get("username","")
    password = data.get("password","")
    user = get_user(username)
    if not user or not verify_password(password, user["password_hash"]):
        return jsonify({"ok": False, "message": "\u7528\u6237\u540d\u6216\u5bc6\u7801\u9519\u8bef"}), 403
    if is_expired(user["expires_at"]):
        return jsonify({"ok": False, "reason": "card_expired", "message": "\u8d26\u53f7\u5df2\u5230\u671f\uff0c\u8bf7\u7eed\u8d39"}), 403
    session["username"] = username
    session["expires_at"] = user["expires_at"]
    return jsonify({
        "ok": True, "message": "\u767b\u5f55\u6210\u529f", "username": username,
        "expires_at": user["expires_at"],
        "countdown": countdown_from_expires(user["expires_at"]),
    })

@app.route("/api/auth/register", methods=["POST"])
async def api_register():
    data = await request.get_json(force=True)
    result = register_user(data.get("card_code",""), data.get("username",""), data.get("password",""), data.get("confirm_password",""))
    if not result.get("ok"):
        return jsonify(result), 403
    return jsonify(result)

@app.route("/api/auth/renew", methods=["POST"])
async def api_renew():
    data = await request.get_json(force=True)
    result = renew_user(data.get("username",""), data.get("password",""), data.get("card_code",""))
    if not result.get("ok"):
        return jsonify(result), 403
    return jsonify(result)

@app.route("/api/logout", methods=["POST"])
async def api_logout():
    logout_user()
    return jsonify({"ok": True})

@app.route("/api/files/download/<path:rel_path>")
async def api_download(rel_path):
    user = current_user()
    if not user:
        return jsonify({"error": "未登录"}), 401
    full_path = os.path.join(os.path.dirname(__file__), "data", "files", rel_path)
    if not os.path.exists(full_path):
        return jsonify({"error": "文件不存在"}), 404
    download_name = request.args.get("name", os.path.basename(rel_path))
    return await send_file(full_path, as_attachment=True, attachment_filename=download_name)

# ===== Pages =====
@app.route("/")
async def index():
    user = current_user()
    if not user:
        return redirect("/login")
    return await render_template("panel.html")

@app.route("/login")
async def login_page():
    return await render_template("login.html")

# ===== WebSocket handler (native async) =====
@app.route("/api/debug/ws_status")
async def debug_ws_status():
    from task_runner import runner
    return jsonify({"ws_clients": len(runner.ws_clients), "running": runner.running, "loop_ok": runner._loop is not None})

@app.websocket("/ws")
async def ws_handler():
    user = None
    username = None
    # Method 1: try Quart session (works if ws context has session)
    try:
        username = session.get("username")
    except:
        pass

    # Method 2: parse cookie from WebSocket headers
    if not username:
        cookie_header = ""
        try:
            cookie_header = websocket.headers.get("cookie", "") or websocket.headers.get("Cookie", "")
        except:
            pass
        if cookie_header:
            for part in cookie_header.split(";"):
                part = part.strip()
                if part.lower().startswith("session="):
                    cookie_value = part.split("=", 1)[1]
                    from itsdangerous import URLSafeTimedSerializer
                    try:
                        s = URLSafeTimedSerializer(app.secret_key)
                        data = s.loads(cookie_value)
                        username = data.get("username", "")
                    except:
                        pass
                    break

    if username:
        from store import get_user, get_users
        users = get_users()
        user = users.get(username)

    if not user:
        await websocket.send(json.dumps({"type":"error","error":"未登录或已过期"}))
        await websocket.close(code=1000)
        return

    import asyncio as _asyncio
    try:
        runner._loop = _asyncio.get_running_loop()
    except:
        pass
    runner.add_client(websocket)

    opts = get_full_options_with_labels()
    auth_info = {
        "username": user["username"],
        "authorized": True,
        "expires_at": user.get("expires_at", ""),
    }
    servers = get_servers()
    bootstrap = {
        "type": "bootstrap",
        "options": opts,
        "auth": auth_info,
        "servers": servers,
        "count": len(servers),
        "source_name": f"累计 {len(servers)} 台" if servers else "",
        "task": {
            "running": runner.running,
            "exit_code": runner.exit_code,
            "task_name": runner.task_name,
            "offset": runner.offset,
            "byte_offset": runner.byte_offset,
            "lines": [],
            "done": False,
        }
    }
    await websocket.send(json.dumps(bootstrap, ensure_ascii=False))

    try:
        while True:
            msg_raw = await websocket.receive()
            if msg_raw is None:
                break
            try:
                msg = json.loads(msg_raw)
            except:
                continue

            mtype = msg.get("type", "")

            if mtype in ("ping", "pong", "subscribe", "resume"):
                continue

            if mtype == "options.get":
                opts = get_full_options_with_labels()
                resp = {"type": "options", "options": opts, "id": msg.get("id")}
                await websocket.send(json.dumps(resp, ensure_ascii=False))

            elif mtype == "options.save":
                body = {k: v for k, v in msg.items() if k not in ("type","id")}
                opts = get_options()
                opts.update(body)
                for k in list(opts.keys()):
                    if k.startswith("type") or k == "id":
                        opts.pop(k, None)
                save_options(opts)
                opts = get_full_options_with_labels()
                resp = {"type": "options", "options": opts, "id": msg.get("id")}
                await websocket.send(json.dumps(resp, ensure_ascii=False))

            elif mtype == "prefs.quick_paste":
                payload = {k: v for k, v in msg.items() if k not in ("type","id")}
                opts = get_options()
                opts["quick_paste"] = payload
                save_options(opts)
                opts = get_full_options_with_labels()
                resp = {"type": "options", "ok": True, "options": opts, "id": msg.get("id")}
                await websocket.send(json.dumps(resp, ensure_ascii=False))

            elif mtype == "servers.list":
                servers = get_servers()
                for s in servers:
                    tok = s.get("token", "")
                    if len(tok) > 14:
                        s["token"] = tok[:8] + "..." + tok[-6:]
                    cred = s.get("cred", "")
                    if cred and len(cred) > 4:
                        s["cred"] = "************"
                resp = {"type": "servers", "servers": servers, "count": len(servers),
                        "source_name": f"累计 {len(servers)} 台" if servers else "", "id": msg.get("id")}
                await websocket.send(json.dumps(resp, ensure_ascii=False))

            elif mtype == "servers.clear":
                clear_servers()
                resp = {"type": "servers", "servers": [], "count": 0, "source_name": "", "id": msg.get("id")}
                await websocket.send(json.dumps(resp, ensure_ascii=False))

            elif mtype == "servers.quick_paste.preview":
                preview = _quick_paste_preview(msg)
                preview["id"] = msg.get("id")
                await websocket.send(json.dumps(preview, ensure_ascii=False))

            elif mtype == "servers.quick_paste.save":
                result = _quick_paste_save(msg)
                result["id"] = msg.get("id")
                await websocket.send(json.dumps(result, ensure_ascii=False))

            elif mtype == "templates.list":
                data = get_templates()
                resp = {"type": "templates", "templates": data["templates"], "active_id": data["active_id"], "id": msg.get("id")}
                await websocket.send(json.dumps(resp, ensure_ascii=False))

            elif mtype == "templates.save":
                name = msg.get("name", "未命名方案")
                opts_snap = get_options()
                stack_id = opts_snap.get("mail_stack", "")
                dkim_id = opts_snap.get("dkim_sign_mode", "")
                count = opts_snap.get("dkim_sign_count", 1)
                from datetime import datetime as _dt
                tid = add_template(name, opts_snap)
                data = get_templates()
                for t in data["templates"]:
                    if not t.get("updated_at"):
                        t["updated_at"] = _dt.now().strftime("%Y-%m-%d %H:%M:%S")
                    if not t.get("summary"):
                        t["summary"] = f"{stack_id} / {dkim_id} / x{count}"
                    if "options" in t and "config" not in t:
                        t["config"] = t.pop("options")
                save_templates(data)
                opts_ret = get_full_options_with_labels()
                resp = {"type": "templates", "templates": data["templates"], "active_id": data["active_id"], "options": opts_ret, "id": msg.get("id")}
                await websocket.send(json.dumps(resp, ensure_ascii=False))

            elif mtype == "templates.apply":
                tid = msg.get("id", "")
                data = get_templates()
                tmpl = next((t for t in data["templates"] if t["id"] == tid), None)
                if tmpl:
                    opts = get_options()
                    tmpl_opts = tmpl.get("config", tmpl.get("options", {}))
                    opts.update(tmpl_opts)
                    save_options(opts)
                    opts = get_full_options_with_labels()
                    set_active_template(tid)
                    resp = {"type": "options", "options": opts, "id": msg.get("id")}
                else:
                    resp = {"type": "error", "error": "方案不存在", "id": msg.get("id")}
                await websocket.send(json.dumps(resp, ensure_ascii=False))

            elif mtype == "templates.delete":
                tid = msg.get("id", "")
                delete_template(tid)
                data = get_templates()
                resp = {"type": "templates", "templates": data["templates"], "active_id": data["active_id"], "id": msg.get("id")}
                await websocket.send(json.dumps(resp, ensure_ascii=False))

            elif mtype == "mailbox.export":
                servers = get_servers()
                rows = []
                for s in servers:
                    if s.get("status") == "成功":
                        rows.append({"email": s.get("email", f"admin@{s['hostname']}"), "password": s.get("cred", ""), "hostname": s.get("hostname", "")})
                resp = {"type": "mailbox", "ok": True, "count": len(rows), "rows": rows, "filename": f"伯乐格式{int(time.time())}.txt", "id": msg.get("id")}
                await websocket.send(json.dumps(resp, ensure_ascii=False))

            elif mtype == "monitor.get":
                resp = {"type": "monitor", "loaded": False, "export_ready": False,
                        "last_pull": {"raw_records":0,"total":0,"delivered":0,"success_rate":0,"hard_bounce":0,"deferred":0,"rejected":0,"soft_bounce":0,"pending":0},
                        "id": msg.get("id")}
                await websocket.send(json.dumps(resp, ensure_ascii=False))

            elif mtype == "monitor.export":
                cat = msg.get("category", "")
                if cat == "full":
                    resp = {"type": "files", "files": [], "error": "没有可导出的数据", "id": msg.get("id")}
                else:
                    resp = {"type": "files", "files": [], "ok": True, "message": "导出完成", "id": msg.get("id")}
                await websocket.send(json.dumps(resp, ensure_ascii=False))

            elif mtype == "files.list":
                resp = {"type": "files", "files": [], "id": msg.get("id")}
                await websocket.send(json.dumps(resp, ensure_ascii=False))

            elif mtype == "task.start":
                action = msg.get("action", "")
                opts = msg.get("options", {})
                confirmed = msg.get("confirmed", False)
                saved_opts = get_options()
                saved_opts.update(opts)
                from options_meta import OPTIONS_META
                # Fill missing defaults, don't overwrite user settings
                for k, v in OPTIONS_META.items():
                    if k not in saved_opts:
                        saved_opts[k] = v
                result = runner.start(action, saved_opts, confirmed)
                resp = {"type": "task.start", "message": result.get("message", ""), "id": msg.get("id")}
                if result.get("error"):
                    resp["error"] = result["error"]
                await websocket.send(json.dumps(resp, ensure_ascii=False))

            elif mtype == "task.stop":
                result = runner.stop()
                resp = {"type": "task.stop", "ok": True, "message": result["message"], "running": result["running"], "id": msg.get("id")}
                await websocket.send(json.dumps(resp, ensure_ascii=False))

            elif mtype == "task.clear_output":
                result = runner.clear_output()
                resp = {"type": "task.clear_output", "ok": True, "offset": result["offset"], "running": False, "task_name": runner.task_name, "id": msg.get("id")}
                await websocket.send(json.dumps(resp, ensure_ascii=False))

            else:
                resp = {"type": "error", "error": "未知消息类型：" + mtype, "id": msg.get("id")}
                await websocket.send(json.dumps(resp, ensure_ascii=False))

    except Exception as e:
        pass
    finally:
        runner.remove_client(websocket)

# ===== Quick paste helpers =====
def _gen_hostname(prefix, start_num, idx, domain):
    if prefix:
        return f"{prefix}{start_num + idx}.{domain}"
    else:
        length = random.randint(4, 8)
        name = "".join(random.choice(string.ascii_lowercase) for _ in range(length))
        return f"{name}.{domain}"

def _quick_paste_preview(msg):
    ip_text = msg.get("ip_text", "")
    domain_text = msg.get("domain_text", "")
    prefix = msg.get("prefix", "")
    start_number = int(msg.get("start_number", 1))
    port = msg.get("port", "22")
    user = msg.get("user", "root")
    password = msg.get("ssh_password", msg.get("password", ""))
    token = msg.get("default_token", "")
    email_prefix = msg.get("email_prefix", "admin")

    cur_opts = get_full_options_with_labels()
    stack_label = cur_opts.get("mail_stack_label", "")

    ips = [ip.strip() for ip in re.split(r'[,\n\s]+', ip_text) if ip.strip()]
    domain_lines = [d.strip() for d in domain_text.split("\n") if d.strip()]
    domains = []
    for line in domain_lines:
        parts = [p.strip() for p in line.split(",") if p.strip()]
        if parts:
            domains.append(parts[0])

    if not ips or not domains:
        return {"type":"servers.quick_paste.preview", "error":"请至少粘贴一个有效 IPv4 地址"} if not ips else {"type":"servers.quick_paste.preview", "error":"请至少粘贴一个主域名"}
    if not token:
        return {"type":"servers.quick_paste.preview", "error":f"以下主域名缺少 Cloudflare Token：{domains[0]}"}

    lines = []
    previews = []
    for idx, ip in enumerate(ips):
        domain = domains[idx % len(domains)]
        hostname = _gen_hostname(prefix, start_number, idx, domain)
        email = f"{email_prefix}@{hostname}"
        line = f"{ip},{port},{user},{password},{domain},{token},{email},{password},{hostname},cloudflare"
        lines.append(line)
        prev = f"{ip} -> {hostname} (admin@{hostname})"
        if stack_label:
            prev += f" [当前栈：{stack_label}]"
        previews.append(prev)

    return {"type":"quick_paste.preview", "count":len(ips), "lines":lines, "previews":previews}

def _quick_paste_save(msg):
    ip_text = msg.get("ip_text", "")
    domain_text = msg.get("domain_text", "")
    prefix = msg.get("prefix", "")
    start_number = int(msg.get("start_number", 1))
    port = msg.get("port", "22")
    user = msg.get("user", "root")
    password = msg.get("ssh_password", msg.get("password", ""))
    token = msg.get("default_token", "")
    email_prefix = msg.get("email_prefix", "admin")

    cur_opts = get_full_options_with_labels()
    stack_label = cur_opts.get("mail_stack_label", "")

    ips = [ip.strip() for ip in re.split(r'[,\n\s]+', ip_text) if ip.strip()]
    domain_lines = [d.strip() for d in domain_text.split("\n") if d.strip()]
    domains = []
    for line in domain_lines:
        parts = [p.strip() for p in line.split(",") if p.strip()]
        if parts:
            domains.append(parts[0])

    if not ips or not domains:
        return {"type":"servers.quick_paste.save", "error":"缺少 IP 或域名", "added":0, "updated":0, "total":0, "count":0, "servers":[]}

    existing = get_servers()
    existing_ips = {s["ip"] for s in existing}
    new_servers = []
    added = 0; updated = 0

    for idx, ip in enumerate(ips):
        domain = domains[idx % len(domains)]
        hostname = _gen_hostname(prefix, start_number, idx, domain)
        email = f"{email_prefix}@{hostname}"
        server = {"id": hostname, "ip": ip, "port": port, "user": user,
                  "cred": password, "hostname": hostname, "domain": domain,
                  "token": token, "email": email,
                  "stack": stack_label, "status": "未处理"}
        if ip in existing_ips:
            for s in existing:
                if s["ip"] == ip:
                    s.update({"hostname": hostname, "domain": domain, "email": email, "token": token, "cred": password, "port": port, "user": user, "stack": stack_label})
                    updated += 1
                    break
        else:
            new_servers.append(server)
            added += 1

    all_servers = existing + new_servers
    save_servers(all_servers)

    return {"type":"servers.quick_paste.save", "added":added, "updated":updated, "total":len(all_servers), "count":len(ips), "servers":all_servers}

if __name__ == "__main__":
    app.run(host="127.0.0.1", port=5000)