import json
import os
import sys
import time
import uuid

import requests
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry
from boto3 import client as s3_client
from botocore.config import Config
from dotenv import load_dotenv

load_dotenv()

APP_ID = os.getenv("X_API_APP_KEY")
ACCESS_TOKEN = os.getenv("X_API_ACCESS_KEY")
RESOURCE_ID = os.getenv("X_API_RESOURCE_ID")
S3_ENDPOINT = os.getenv("S3_ENDPOINT")
S3_ACCESS_KEY_ID = os.getenv("S3_ACCESS_KEY_ID")
S3_SECRET_ACCESS_KEY = os.getenv("S3_SECRET_ACCESS_KEY")
S3_BUCKET = os.getenv("S3_BUCKET")

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

retry = Retry(total=5, backoff_factor=1, status_forcelist=[500, 502, 503, 504])
session = requests.Session()
session.mount("http://", HTTPAdapter(max_retries=retry))
session.mount("https://", HTTPAdapter(max_retries=retry))

s3 = s3_client(
    "s3",
    endpoint_url=S3_ENDPOINT,
    aws_access_key_id=S3_ACCESS_KEY_ID,
    aws_secret_access_key=S3_SECRET_ACCESS_KEY,
    config=Config(signature_version="s3v4"),
    region_name="auto",
)


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)


def make_headers(task_id, with_sequence=False):
    headers = {
        "X-Api-App-Key": APP_ID,
        "X-Api-Access-Key": ACCESS_TOKEN,
        "X-Api-Resource-Id": RESOURCE_ID,
        "X-Api-Request-Id": task_id,
    }
    if with_sequence:
        headers["X-Api-Sequence"] = "-1"
    return headers


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} ...")
s3.upload_file(audio_file, S3_BUCKET, obj_name)
url = s3.generate_presigned_url(
    "get_object",
    Params={"Bucket": S3_BUCKET, "Key": obj_name},
    ExpiresIn=3600,
)

task_id = str(uuid.uuid4())
headers = make_headers(task_id, with_sequence=True)
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 = session.post(SUBMIT_URL, json=body, headers=headers, timeout=30)
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 = session.post(
            QUERY_URL,
            json={},
            headers={
                **make_headers(task_id),
                "X-Tt-Logid": logid,
            },
            timeout=30,
        )
        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", [])

            os.makedirs("output", exist_ok=True)
            json_path = f"output/{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"output/{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.")