"""
server.py — HarmoSplit バックエンド
Hugging Face Spaces 対応（ポート 7860）+ Stripe 月額決済
"""

import os
import sys
import uuid
import json
import shutil
import secrets
import tempfile
import threading
import subprocess
from pathlib import Path
from datetime import datetime, timezone

from flask import (
    Flask, request, jsonify, send_file,
    Response, send_from_directory, redirect, url_for
)
from flask_cors import CORS

sys.path.insert(0, str(Path(__file__).parent))
import app as core

# ── Flask ──────────────────────────────────────────────────────
flask_app = Flask(__name__, static_folder="static", static_url_path="")
CORS(flask_app)

# ── 設定 ──────────────────────────────────────────────────────
PORT = int(os.environ.get("PORT", 7860))

# Stripe キー（HF Spaces Secrets または .env で設定）
STRIPE_SECRET_KEY    = os.environ.get("STRIPE_SECRET_KEY", "")
STRIPE_WEBHOOK_SECRET = os.environ.get("STRIPE_WEBHOOK_SECRET", "")
STRIPE_PRICE_ID      = os.environ.get("STRIPE_PRICE_ID", "")
# アプリの公開 URL（Webhook / Checkout success URL 用）
APP_URL = os.environ.get("APP_URL", f"http://localhost:{PORT}")

# 決済不要モード（Stripe キー未設定なら無料開放）
FREE_MODE = not bool(STRIPE_SECRET_KEY)

# ── 永続ストレージ ──────────────────────────────────────────
# HF Spaces では /data が永続。ローカルでは ./data を使用。
DATA_DIR = Path("/data") if Path("/data").exists() else Path("./data")
DATA_DIR.mkdir(parents=True, exist_ok=True)

TOKENS_FILE  = DATA_DIR / "tokens.json"
UPLOAD_DIR   = DATA_DIR / "uploads"
UPLOAD_DIR.mkdir(parents=True, exist_ok=True)

# ── トークン管理 ──────────────────────────────────────────────
tokens_lock = threading.Lock()

def load_tokens() -> dict:
    with tokens_lock:
        if TOKENS_FILE.exists():
            try:
                return json.loads(TOKENS_FILE.read_text("utf-8"))
            except Exception:
                pass
        return {}

def save_tokens(data: dict):
    with tokens_lock:
        TOKENS_FILE.write_text(json.dumps(data, indent=2, ensure_ascii=False), "utf-8")

def create_token(customer_id: str, subscription_id: str, email: str) -> str:
    token = secrets.token_urlsafe(32)
    data = load_tokens()
    data[token] = {
        "customer_id":      customer_id,
        "subscription_id":  subscription_id,
        "email":            email,
        "created_at":       datetime.now(timezone.utc).isoformat(),
        "active":           True,
    }
    save_tokens(data)
    return token

def is_token_valid(token: str) -> bool:
    if FREE_MODE:
        return True   # 無料モードは常に有効
    data = load_tokens()
    entry = data.get(token)
    return bool(entry and entry.get("active"))

def deactivate_token_by_subscription(subscription_id: str):
    data = load_tokens()
    for info in data.values():
        if info.get("subscription_id") == subscription_id:
            info["active"] = False
    save_tokens(data)

# ── ジョブ管理 ──────────────────────────────────────────────
JOBS: dict[str, dict] = {}
JOBS_LOCK = threading.Lock()

def log_progress(job_id: str, message: str, percent: int | None = None):
    entry = {"msg": message}
    if percent is not None:
        entry["pct"] = percent
    with JOBS_LOCK:
        if job_id in JOBS:
            JOBS[job_id]["progress"].append(entry)

def _load_stems_helper(demucs_out: Path, model: str, wav_path: Path):
    import soundfile as sf
    import numpy as np
    stem_dir = demucs_out / model / wav_path.stem
    if not stem_dir.exists():
        candidates = list(demucs_out.rglob("*.wav"))
        if not candidates:
            raise FileNotFoundError("ステムファイルが見つかりません")
        stem_dir = candidates[0].parent
    stems = {}
    sr = 44100
    for wav_file in sorted(stem_dir.glob("*.wav")):
        data, sr = sf.read(str(wav_file), always_2d=True)
        stems[wav_file.stem] = data.astype(np.float_())
    return stems, sr

core._load_stems = _load_stems_helper

def process_job(job_id: str, input_path: Path, inst_vol: float, model: str, use_mdx: bool):
    tmp_dir    = Path(tempfile.mkdtemp(prefix=f"hmsplit_{job_id[:8]}_"))
    demucs_out = tmp_dir / "demucs_out"
    demucs_out.mkdir(parents=True, exist_ok=True)
    try:
        with JOBS_LOCK:
            JOBS[job_id]["status"] = "processing"

        log_progress(job_id, "🎵 音声を読み込み中...", 5)
        wav_path = core.prepare_audio(input_path, tmp_dir)

        log_progress(job_id, "🤖 Demucs AI で音源を分離中...", 15)
        cmd = [sys.executable, "-m", "demucs", "-n", model,
               "-o", str(demucs_out), str(wav_path)]
        proc = subprocess.Popen(
            cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
            text=True, encoding="utf-8", errors="replace"
        )
        for line in proc.stdout:
            line = line.rstrip()
            if line.strip():
                log_progress(job_id, line)
        proc.wait()
        if proc.returncode != 0:
            raise RuntimeError("Demucs 処理に失敗しました")

        log_progress(job_id, "📂 ステムを読み込み中...", 55)
        stems, sr = core._load_stems(demucs_out, model, wav_path)

        if model == "htdemucs_6s" and "vocals" not in stems:
            log_progress(job_id, "⚠️ htdemucs にフォールバック", 57)
            shutil.rmtree(demucs_out); demucs_out.mkdir(parents=True, exist_ok=True)
            model = "htdemucs"
            subprocess.run([sys.executable, "-m", "demucs", "-n", model,
                            "-o", str(demucs_out), str(wav_path)], check=True, capture_output=True)
            stems, sr = core._load_stems(demucs_out, model, wav_path)

        mdx_model_path = None
        if use_mdx:
            log_progress(job_id, "🧠 UVR MDX-NET モデルを準備中...", 60)
            try:
                mdx_cache = DATA_DIR / "models"
                mdx_cache.mkdir(parents=True, exist_ok=True)
                mdx_model_path = core.download_mdx_model(mdx_cache)
                log_progress(job_id, "✅ UVR MDX-NET 準備完了", 65)
            except Exception as e:
                log_progress(job_id, f"⚠️ MDX 取得失敗、Mid/Side にフォールバック: {e}", 65)

        log_progress(job_id, "🎚️ L/R パンニング & ミックス処理中...", 68)
        if mdx_model_path:
            log_progress(job_id, "🔬 リード / バッキング AI 分離中... (数分かかります)", 70)
        mixed = core.mix_stems(stems, model, inst_vol, mdx_model_path=mdx_model_path, sr=sr)

        log_progress(job_id, "💾 WAV を書き出し中...", 95)
        import soundfile as sf
        output_path = UPLOAD_DIR / f"{job_id}_panned.wav"
        sf.write(str(output_path), mixed, sr, subtype="PCM_16")

        with JOBS_LOCK:
            JOBS[job_id]["status"]      = "done"
            JOBS[job_id]["output_path"] = str(output_path)
        log_progress(job_id, "✅ 処理完了！ダウンロードボタンをクリックしてください。", 100)

    except Exception as e:
        import traceback
        with JOBS_LOCK:
            JOBS[job_id]["status"] = "error"
            JOBS[job_id]["error"]  = str(e)
        log_progress(job_id, f"❌ {e}")
    finally:
        shutil.rmtree(tmp_dir, ignore_errors=True)
        try:
            input_path.unlink(missing_ok=True)
        except Exception:
            pass

# ── ルーティング ────────────────────────────────────────────

@flask_app.route("/")
def index():
    return send_from_directory("static", "index.html")

@flask_app.route("/pricing")
def pricing():
    return send_from_directory("static", "pricing.html")

@flask_app.route("/success")
def success():
    return send_from_directory("static", "success.html")

@flask_app.route("/legal")
def legal():
    return send_from_directory("static", "legal.html")

@flask_app.route("/auth-mode")
def auth_mode():
    """フロントエンドが無料/有料モードを識別するためのエンドポイント"""
    return jsonify({"free_mode": FREE_MODE})

@flask_app.route("/pricing-info")
def pricing_info():
    """料金ページ用: Stripe から Price 情報を取得して返す"""
    if FREE_MODE:
        return jsonify({"price": 0, "currency": "jpy", "free_mode": True})
    try:
        import stripe
        stripe.api_key = STRIPE_SECRET_KEY
        price = stripe.Price.retrieve(STRIPE_PRICE_ID)
        return jsonify({
            "price":     price.unit_amount,
            "currency":  price.currency,
            "interval":  price.recurring.interval if price.recurring else "month",
        })
    except Exception as e:
        return jsonify({"error": str(e)}), 500

# ── Stripe: Checkout セッション作成 ──────────────────────────
@flask_app.route("/create-checkout", methods=["POST"])
def create_checkout():
    if FREE_MODE:
        return jsonify({"error": "Stripe 未設定（開発モード）"}), 400
    try:
        import stripe
        stripe.api_key = STRIPE_SECRET_KEY
        session = stripe.checkout.Session.create(
            payment_method_types=["card"],
            line_items=[{"price": STRIPE_PRICE_ID, "quantity": 1}],
            mode="subscription",
            success_url=f"{APP_URL}/success?session_id={{CHECKOUT_SESSION_ID}}",
            cancel_url=f"{APP_URL}/pricing",
        )
        return jsonify({"url": session.url})
    except Exception as e:
        return jsonify({"error": str(e)}), 500

# ── Stripe: Webhook ──────────────────────────────────────────
@flask_app.route("/webhook", methods=["POST"])
def stripe_webhook():
    if FREE_MODE:
        return "", 200
    try:
        import stripe
        stripe.api_key = STRIPE_SECRET_KEY
        payload    = request.get_data()
        sig_header = request.headers.get("Stripe-Signature", "")
        event = stripe.Webhook.construct_event(payload, sig_header, STRIPE_WEBHOOK_SECRET)
    except Exception as e:
        return jsonify({"error": str(e)}), 400

    etype = event["type"]
    obj   = event["data"]["object"]

    if etype == "checkout.session.completed":
        sub_id   = obj.get("subscription")
        cust_id  = obj.get("customer")
        email    = obj.get("customer_email") or obj.get("customer_details", {}).get("email", "")
        token    = create_token(cust_id, sub_id, email)
        print(f"[WEBHOOK] 新サブスクリプション: {email} → token={token[:8]}...")

    elif etype in ("customer.subscription.deleted", "customer.subscription.paused"):
        sub_id = obj.get("id")
        deactivate_token_by_subscription(sub_id)
        print(f"[WEBHOOK] サブスクリプション停止: {sub_id}")

    return "", 200

# ── Stripe: 成功後トークン取得 ────────────────────────────────
@flask_app.route("/get-token")
def get_token():
    """決済完了後に Stripe Session ID からトークンを返す"""
    if FREE_MODE:
        return jsonify({"token": "FREE_MODE"})
    session_id = request.args.get("session_id", "")
    if not session_id:
        return jsonify({"error": "session_id が必要です"}), 400
    try:
        import stripe
        stripe.api_key = STRIPE_SECRET_KEY
        session = stripe.checkout.Session.retrieve(session_id)
        sub_id  = session.get("subscription")
        cust_id = session.get("customer")
        email   = session.get("customer_details", {}).get("email", "")
        # 既存トークンを探す（Webhook が先に処理している場合）
        data = load_tokens()
        for tok, info in data.items():
            if info.get("subscription_id") == sub_id:
                return jsonify({"token": tok, "email": email})
        # Webhook がまだなら作成
        token = create_token(cust_id, sub_id, email)
        return jsonify({"token": token, "email": email})
    except Exception as e:
        return jsonify({"error": str(e)}), 500

# ── トークン検証 ──────────────────────────────────────────────
@flask_app.route("/verify-token", methods=["POST"])
def verify_token():
    token = request.json.get("token", "")
    return jsonify({"valid": is_token_valid(token)})

# ── ファイルアップロード & 処理 ────────────────────────────────
@flask_app.route("/upload", methods=["POST"])
def upload():
    # トークン認証
    token = request.form.get("token", "")
    if not is_token_valid(token):
        return jsonify({"error": "無効なトークンです。料金ページから登録してください。"}), 401

    if "file" not in request.files:
        return jsonify({"error": "ファイルがありません"}), 400
    file = request.files["file"]
    if not file.filename:
        return jsonify({"error": "ファイル名が空です"}), 400

    inst_vol = float(request.form.get("inst_vol", 0.15))
    model    = request.form.get("model", "htdemucs_6s")
    use_mdx  = request.form.get("use_mdx", "true").lower() == "true"

    job_id     = str(uuid.uuid4())
    suffix     = Path(file.filename).suffix
    input_path = UPLOAD_DIR / f"{job_id}_input{suffix}"
    file.save(str(input_path))

    with JOBS_LOCK:
        JOBS[job_id] = {
            "status": "queued", "progress": [],
            "output_path": None, "error": None,
            "filename": file.filename,
        }

    t = threading.Thread(
        target=process_job,
        args=(job_id, input_path, inst_vol, model, use_mdx),
        daemon=True,
    )
    t.start()
    return jsonify({"job_id": job_id})

# ── 進捗 SSE ──────────────────────────────────────────────────
@flask_app.route("/progress/<job_id>")
def progress(job_id: str):
    def generate():
        import time
        sent = 0
        while True:
            with JOBS_LOCK:
                if job_id not in JOBS:
                    yield 'data: {"error":"not found"}\n\n'; return
                job      = JOBS[job_id]
                new_ents = job["progress"][sent:]
                sent    += len(new_ents)
                status   = job["status"]
            for e in new_ents:
                yield f"data: {json.dumps(e, ensure_ascii=False)}\n\n"
            if status in ("done", "error"):
                yield f"data: {json.dumps({'status': status})}\n\n"; return
            time.sleep(0.5)
    return Response(generate(), mimetype="text/event-stream",
                    headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"})

@flask_app.route("/status/<job_id>")
def status(job_id: str):
    with JOBS_LOCK:
        job = JOBS.get(job_id)
    if not job:
        return jsonify({"error": "not found"}), 404
    return jsonify({"status": job["status"], "error": job.get("error")})

# ── ダウンロード ──────────────────────────────────────────────
@flask_app.route("/download/<job_id>")
def download(job_id: str):
    with JOBS_LOCK:
        job = JOBS.get(job_id)
    if not job or job["status"] != "done":
        return jsonify({"error": "ファイル未準備"}), 404
    output_path = Path(job["output_path"])
    if not output_path.exists():
        return jsonify({"error": "ファイルが見つかりません"}), 404
    download_name = Path(job.get("filename", "audio")).stem + "_panned.wav"
    return send_file(str(output_path), as_attachment=True,
                     download_name=download_name, mimetype="audio/wav")

# ── 起動 ──────────────────────────────────────────────────────
if __name__ == "__main__":
    mode = "FREE（Stripe 未設定）" if FREE_MODE else "有料（Stripe 有効）"
    print("=" * 52)
    print(f"🎵 HarmoSplit 起動中... モード: {mode}")
    print(f"   http://localhost:{PORT}")
    print("=" * 52)
    flask_app.run(host="0.0.0.0", port=PORT, debug=False, threaded=True)