"""
flash_asr.py — AnyFile2MD 音频转写模块 v2.0

架构修正（2026-04-15）：
  ❌ 旧方案：base64 multipart → DashScope 同步 API（10MB 限制）
  ✅ 新方案：OSS 签名 URL → DashScope paraformer-v2 异步任务（无大小限制）

职责（封闭函数，铁律 1）：
  transcribe_from_oss_url(url) → AsyncIterator[{"text", "begin_ms", "end_ms"}]
  generate_oss_presign(user_id, ext) → {"upload_url", "oss_key", "read_url"}
  is_supported_audio(filename) → bool

不调用 LLM，不写 DB，不持久化任何状态。
"""

import asyncio
import hmac
import hashlib
import base64
import os
import time
import logging
from datetime import datetime, timezone
from typing import AsyncIterator
from urllib.parse import quote

import httpx

logger = logging.getLogger(__name__)

# ─── 支持格式 ────────────────────────────────────────────────
SUPPORTED_EXTENSIONS = {
    ".mp3", ".mp4", ".m4a", ".wav", ".webm",
    ".ogg", ".aac", ".flac", ".opus", ".amr",
}

def is_supported_audio(filename: str) -> bool:
    ext = os.path.splitext(filename)[1].lower()
    return ext in SUPPORTED_EXTENSIONS


# ─── OSS 配置 ─────────────────────────────────────────────────
def _oss_config() -> dict:
    return {
        "access_key_id":     os.getenv("OSS_ACCESS_KEY_ID", ""),
        "access_key_secret": os.getenv("OSS_ACCESS_KEY_SECRET", ""),
        "bucket":            os.getenv("OSS_BUCKET", "meetings-dev"),
        "endpoint":          os.getenv("OSS_ENDPOINT", "oss-cn-guangzhou.aliyuncs.com"),
        "prefix":            os.getenv("OSS_AUDIO_PREFIX", "mindos-next/audio"),
    }


def generate_oss_presign(user_id: str, ext: str, expire_seconds: int = 3600,
                         prefix: str | None = None) -> dict:
    """
    生成 OSS 预签名 URL（使用 oss2 SDK，V2 验证过的可靠方式）。
    返回：{ upload_url, oss_key, read_url, content_type }
      - upload_url ：前端直传用（HTTP PUT，必须带 Content-Type: application/octet-stream）
      - read_url   ：交给 DashScope paraformer-v2 下载用
      - content_type：前端 PUT 时必须传此 Content-Type

    Args:
        prefix: OSS 路径前缀。默认 None 时使用 OSS_AUDIO_PREFIX 环境变量。
                文档管线传 "mindos-next/docs"，音频管线不传（使用默认值）。
    """
    import oss2  # type: ignore

    cfg = _oss_config()
    if not cfg["access_key_id"] or not cfg["access_key_secret"]:
        raise RuntimeError("[FlashASR] OSS_ACCESS_KEY_ID / OSS_ACCESS_KEY_SECRET 未配置")

    actual_prefix = prefix or cfg["prefix"]
    ts = int(time.time())
    oss_key = f"{actual_prefix}/{user_id}/{ts}_{user_id[:8]}{ext}"

    auth   = oss2.Auth(cfg["access_key_id"], cfg["access_key_secret"])
    bucket = oss2.Bucket(auth, f"https://{cfg['endpoint']}", cfg["bucket"])

    # Content-Type 必须写入签名，前端 PUT 时须传相同值
    content_type = "application/octet-stream"
    upload_url = bucket.sign_url(
        "PUT", oss_key, expire_seconds,
        slash_safe=True,
        headers={"Content-Type": content_type},
    )
    read_url = bucket.sign_url("GET", oss_key, expire_seconds, slash_safe=True)

    host = f"{cfg['bucket']}.{cfg['endpoint']}"
    logger.info("[FlashASR] presign ok oss_key=%s", oss_key)

    return {
        "upload_url":   upload_url,
        "read_url":     read_url,
        "oss_key":      oss_key,
        "host":         host,
        "content_type": content_type,
    }


# ─── DashScope paraformer-v2 异步任务 ─────────────────────────
DASHSCOPE_API_KEY = lambda: os.getenv("DASHSCOPE_API_KEY", "")
DASHSCOPE_HOST    = "https://dashscope.aliyuncs.com"
TRANSCRIPTION_URL = f"{DASHSCOPE_HOST}/api/v1/services/audio/asr/transcription"

_HEADERS = lambda: {
    "Authorization": f"Bearer {DASHSCOPE_API_KEY()}",
    "Content-Type":  "application/json",
    "X-DashScope-Async": "enable",
}


async def _submit_task(file_url: str) -> str:
    """提交 paraformer-v2 异步转写任务，返回 task_id。"""
    body = {
        "model": "paraformer-v2",
        "input": {"file_urls": [file_url]},
        "parameters": {
            "language_hints": ["zh", "en"],
            "timestamp_alignment": True,
        },
    }
    async with httpx.AsyncClient(timeout=30) as client:
        resp = await client.post(TRANSCRIPTION_URL, json=body, headers=_HEADERS())
        resp.raise_for_status()
        data = resp.json()

    task_id = data.get("output", {}).get("task_id")
    if not task_id:
        raise RuntimeError(f"[FlashASR] 提交任务失败: {data}")
    logger.info("[FlashASR] 任务已提交 task_id=%s", task_id)
    return task_id


async def _poll_task(task_id: str, poll_interval: float = 5.0, timeout: float = 600.0) -> dict:
    """轮询任务状态，返回完成后的 output。"""
    url = f"{DASHSCOPE_HOST}/api/v1/tasks/{task_id}"
    deadline = time.monotonic() + timeout

    async with httpx.AsyncClient(timeout=30) as client:
        while time.monotonic() < deadline:
            resp = await client.get(url, headers=_HEADERS())
            resp.raise_for_status()
            data = resp.json()
            status = data.get("output", {}).get("task_status", "")

            if status == "SUCCEEDED":
                logger.info("[FlashASR] 任务完成 task_id=%s", task_id)
                return data["output"]
            elif status in ("FAILED", "CANCELED"):
                raise RuntimeError(f"[FlashASR] 任务失败 task_id={task_id} status={status}: {data}")

            logger.debug("[FlashASR] 轮询中 task_id=%s status=%s", task_id, status)
            await asyncio.sleep(poll_interval)

    raise RuntimeError(f"[FlashASR] 任务超时 task_id={task_id} (>{timeout}s)")


async def transcribe_from_oss_url(read_url: str) -> AsyncIterator[dict]:
    """
    调用 DashScope paraformer-v2 异步转写。
    yield {"text": str, "begin_ms": int, "end_ms": int}
    """
    if not DASHSCOPE_API_KEY():
        raise RuntimeError("[FlashASR] DASHSCOPE_API_KEY 未配置")

    task_id = await _submit_task(read_url)
    output  = await _poll_task(task_id)

    # 下载转写结果 JSON
    result_url = output.get("results", [{}])[0].get("transcription_url", "")
    if not result_url:
        raise RuntimeError(f"[FlashASR] 未找到 transcription_url: {output}")

    async with httpx.AsyncClient(timeout=60) as client:
        resp = await client.get(result_url)
        resp.raise_for_status()
        transcript_data = resp.json()

    # 解析句子列表
    sentences = transcript_data.get("transcripts", [{}])[0].get("sentences", [])
    if not sentences:
        # fallback：整段文本
        text = transcript_data.get("transcripts", [{}])[0].get("text", "")
        if text:
            yield {"text": text, "begin_ms": 0, "end_ms": 0}
        return

    for s in sentences:
        yield {
            "text":     s.get("text", "").strip(),
            "begin_ms": int(s.get("begin_time", 0)),
            "end_ms":   int(s.get("end_time", 0)),
        }
