import json
import os
import sys
import time
import uuid
from datetime import timedelta

import requests
from dotenv import load_dotenv
from minio import Minio

load_dotenv()

API_KEY = os.getenv("X_API_KEY")
RESOURCE_ID = os.getenv("X_API_RESOURCE_ID")
MINIO_ENDPOINT = os.getenv("MINIO_ENDPOINT")
MINIO_ACCESS_KEY = os.getenv("MINIO_ACCESS_KEY")
MINIO_SECRET_KEY = os.getenv("MINIO_SECRET_KEY")
MINIO_BUCKET = os.getenv("MINIO_BUCKET")

SUBMIT_URL = "https://openspeech.bytedance.com/api/v3/auc/bigmodel/submit"
QUERY_URL = "https://openspeech.bytedance.com/api/v3/auc/bigmodel/query"

client = Minio(
    MINIO_ENDPOINT,
    access_key=MINIO_ACCESS_KEY,
    secret_key=MINIO_SECRET_KEY,
    secure=True,
)

def fmt_time(ms):
    s = ms / 1000
    m = int(s // 60)
    sec = s % 60
    return f"{m:02d}:{sec:05.2f}"


def build_markdown(data, audio_file):
    result = data.get("result", {})
    utterances = result.get("utterances", [])
    duration = data.get("audio_info", {}).get("duration", 0)

    speaker_order = []
    for u in utterances:
        sp = u.get("additions", {}).get("speaker", "unknown")
        if sp not in speaker_order:
            speaker_order.append(sp)

    lines = [
        "# 转写结果",
        "",
        f"**源文件**: {audio_file}  ",
        f"**音频时长**: {fmt_time(duration)}  ",
        f"**说话人数量**: {len(speaker_order)} 人  ",
        f"**分句数量**: {len(utterances)}",
        "",
        "---",
        "",
        "## 按时间线",
        "",
    ]
    for u in utterances:
        sp = u.get("additions", {}).get("speaker", "?")
        start = u.get("start_time", 0)
        end = u.get("end_time", 0)
        t = u.get("text", "").strip()
        if not t:
            continue
        lines.append(f"- **{fmt_time(start)} - {fmt_time(end)}** **说话人{sp}** {t}")

    lines.append("")
    return "\n".join(lines)


audio_file = sys.argv[1] if len(sys.argv) > 1 else "2026年05月26日 17点53分.mp3"
audio_ext = os.path.splitext(audio_file)[1].lstrip(".")
audio_name = os.path.splitext(os.path.basename(audio_file))[0]

obj_name = f"audio/{uuid.uuid4().hex[:8]}_{os.path.basename(audio_file)}"
print(f"Uploading {audio_file} ...")
client.fput_object(MINIO_BUCKET, obj_name, audio_file)
url = client.presigned_get_object(MINIO_BUCKET, obj_name, expires=timedelta(hours=1))

task_id = str(uuid.uuid4())
headers = {
    "X-Api-Key": API_KEY,
    "X-Api-Resource-Id": RESOURCE_ID,
    "X-Api-Request-Id": task_id,
    "X-Api-Sequence": "-1",
}
body = {
    "user": {"uid": "audioscribe"},
    "audio": {"url": url, "format": audio_ext},
    "request": {
        "model_name": "bigmodel",
        "enable_itn": True,
        "enable_punc": True,
        "enable_ddc": True,
        "enable_speaker_info": True,
        "show_utterances": True,
    },
}

resp = requests.post(SUBMIT_URL, json=body, headers=headers)
logid = resp.headers.get("X-Tt-Logid", "")
status = resp.headers.get("X-Api-Status-Code", "")
print(f"Submit: {status}")

if status != "20000000":
    print(f"FAILED: {resp.headers.get('X-Api-Message', '')}")
else:
    while True:
        resp = requests.post(
            QUERY_URL,
            json={},
            headers={
                "X-Api-Key": API_KEY,
                "X-Api-Resource-Id": RESOURCE_ID,
                "X-Api-Request-Id": task_id,
                "X-Tt-Logid": logid,
            },
        )
        code = resp.headers.get("X-Api-Status-Code", "")
        if code == "20000000":
            data = resp.json()
            text = data.get("result", {}).get("text", "")
            utterances = data.get("result", {}).get("utterances", [])

            json_path = f"{audio_name}_转写结果.json"
            with open(json_path, "w", encoding="utf-8") as f:
                json.dump(data, f, ensure_ascii=False, indent=2)

            md_path = f"{audio_name}_转写结果.md"
            with open(md_path, "w", encoding="utf-8") as f:
                f.write(build_markdown(data, audio_file))

            print(f"Done ({len(text)} chars, {len(utterances)} utterances)")
            print(f"Saved: {json_path}, {md_path}")
            break
        elif code in ("20000001", "20000002"):
            print(f"  Processing... ({code})")
            time.sleep(5)
        else:
            print(f"FAILED: {code} {resp.headers.get('X-Api-Message', '')}")
            break

print("\nAll done.")