import os
from io import BytesIO
import re
import requests
import base64
import tempfile
from pathlib import Path
from typing import Union
import httpx
import subprocess

from langchain_core.messages import HumanMessage
from langchain_openai import ChatOpenAI


class QwenAsr:
    def __init__(self, base_url: str, model: str, temperature: float = 0.0, api_key: str = 'no_key'):
        self.base_url = base_url
        self.model = model
        self.temperature = temperature
        self.api_key = api_key
        # self.asr_model = ChatOpenAI(base_url="http://192.168.1.59:8101/v1", model="Qwen/Qwen3-ASR-1.7B", temperature=0.7,
        #                        api_key='no_key')
        self.asr_model = ChatOpenAI(base_url=self.base_url, model=self.model, temperature=self.temperature,
                                    api_key=self.api_key)

    @staticmethod
    def _audio_content(source: Union[str, Path]) -> list:
        """
        读取音频/视频文件并构建 LLM 消息内容。
        如果 source 是视频，会先提取音频；如果是 HTTP 音频 URL，直接引用 URL。

        Args:
            source: 音频/视频文件路径或 URL
        Returns:
            list: LLM 消息内容列表
        """
        source_str = str(source)
        audio_exts = {"wav", "mp3", "m4a", "flac", "ogg", "webm", "aac", "wma"}
        video_exts = {"mp4", "avi", "mkv", "mov", "wmv", "flv", "webm", "mpeg", "mpg", "ts", "3gp", "m4v"}
        temp_files = []

        try:
            content = []
            is_url = source_str.startswith(("http://", "https://"))
            ext = Path(source_str).suffix.lstrip(".").lower()

            # HTTP 音频 URL：直接引用，无需下载
            if is_url and ext in audio_exts:
                # audio_prompt = f"以下是音频文件({audio_name})的内容，请进行分析。"
                # content = [{"type": "text", "text": audio_prompt}]
                # if prompt is not None:
                #     content.append({"type": "text", "text": prompt})
                content.append({
                    "type": "audio_url",
                    "audio_url": {
                        "url": source_str
                    }
                })
                return content

            # 本地文件或需要下载的 URL（视频或其他媒体）
            if is_url:
                resp = httpx.get(source_str, timeout=60.0)
                resp.raise_for_status()
                suffix = Path(source_str).suffix or ".mp4"
                temp_file = tempfile.NamedTemporaryFile(suffix=suffix, delete=False)
                temp_file.write(resp.content)
                temp_file.close()
                temp_files.append(temp_file.name)
                media_path = temp_file.name
            else:
                media_path = str(source)
                if not Path(media_path).exists():
                    raise FileNotFoundError(f"文件不存在: {media_path}")

            media_ext = Path(media_path).suffix.lstrip(".").lower()

            # 如果是视频，提取音频为 wav
            if media_ext in video_exts or (is_url and ext not in audio_exts):
                temp_audio = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
                temp_audio.close()
                temp_files.append(temp_audio.name)
                cmd = [
                    "ffmpeg",
                    "-y",
                    "-i", media_path,
                    "-vn",
                    "-acodec", "pcm_s16le",
                    "-ar", "16000",
                    "-ac", "1",
                    "-hide_banner",
                    "-loglevel", "error",
                    temp_audio.name
                ]
                result = subprocess.run(cmd, capture_output=True, text=True)
                if result.returncode != 0:
                    raise RuntimeError(f"ffmpeg 提取音频失败: {result.stderr.strip()}")
                audio_path = temp_audio.name
                media_ext = "wav"
            else:
                audio_path = media_path
                if media_ext not in audio_exts:
                    media_ext = "wav"

            # 读取音频并 base64 编码
            with open(audio_path, "rb") as f:
                audio_bytes = f.read()
            audio_base64 = base64.b64encode(audio_bytes).decode("utf-8")

            # 构建消息内容
            # audio_prompt = f"以下是音频文件({audio_name})的内容，请进行分析。"
            # content = [{"type": "text", "text": audio_prompt}]
            # if prompt is not None:
            #     content.append({"type": "text", "text": prompt})

            content.append({
                "type": "audio_url",
                "audio_url": {
                    "url": f"data:audio/{media_ext};base64,{audio_base64}"
                }
            })

            return content
        finally:
            # 清理所有临时文件
            for tf in temp_files:
                p = Path(tf)
                if p.exists():
                    p.unlink()

    @staticmethod
    def _extract_asr_text(text: str) -> str:
        """
        从ASR响应文本中提取<asr_text>标签内容

        Args:
            text: 原始响应文本，如 'language Chinese<asr_text>放大一点一倍。'

        Returns:
            str: 提取的转录文本，如 '放大一点一倍。'
        """
        # 尝试匹配 <asr_text>...</asr_text>
        match = re.search(r'<asr_text>(.*?)(?:</asr_text>|$)', text, re.DOTALL)
        if match:
            return match.group(1).strip()
        return text

    def asr(self, source, cache_path:str=None):
        if cache_path is not None and os.path.exists(cache_path):
            with open(cache_path, 'r', encoding='utf-8') as f:
                return f.read()
        contents = self._audio_content(source)
        message = HumanMessage(content=contents)
        result = self.asr_model.invoke([message])
        asr_text = self._extract_asr_text(result.content)
        if cache_path is not None:
            os.makedirs(Path(cache_path).parent, exist_ok=True)
            with open(cache_path, 'w', encoding='utf-8') as f:
                f.write(asr_text)
        return asr_text