get_wechat/chatlog_fastAPI/services/media_resolver.py

from __future__ import annotations

import logging
import sqlite3
from dataclasses import dataclass
from pathlib import Path

import httpx
from fastapi import HTTPException

from config import settings
from services.chatlog_context import get_chatlog_context
from services.http_client import shared_client

log = logging.getLogger(__name__)


@dataclass
class ResolvedMedia:
    bytes: bytes
    content_type: str
    url: str


def _media_url(kind: str, key: str, thumb: bool = False) -> str:
    url = f"{settings.chatlog_base_url}/{kind}/{key}"
    if thumb:
        url += "?thumb=1"
    return url


def _read_voice_resource_status(key: str) -> dict:
    ctx = get_chatlog_context()
    work_dir = ctx.get("work_dir") or ""
    if not work_dir:
        return {"checked": False, "reason": "missing_work_dir"}

    db_path = Path(work_dir) / "db_storage" / "message" / "message_resource.db"
    if not db_path.exists():
        return {"checked": False, "reason": "message_resource_db_missing", "path": str(db_path)}

    try:
        conn = sqlite3.connect(f"file:{db_path.as_posix()}?mode=ro", uri=True)
        conn.row_factory = sqlite3.Row
        try:
            info = conn.execute(
                "SELECT * FROM MessageResourceInfo WHERE message_svr_id=?",
                (int(key),),
            ).fetchone()
            if not info:
                return {
                    "checked": True,
                    "found": False,
                    "path": str(db_path),
                    "message": "当前已解密资源库里没有这条语音的媒体资源记录",
                }
            details = conn.execute(
                "SELECT type,size,status,data_index FROM MessageResourceDetail WHERE message_id=?",
                (info["message_id"],),
            ).fetchall()
            return {
                "checked": True,
                "found": True,
                "path": str(db_path),
                "message_id": info["message_id"],
                "resources": [dict(row) for row in details],
            }
        finally:
            conn.close()
    except Exception as exc:
        return {"checked": False, "reason": "resource_db_read_failed", "error": str(exc), "path": str(db_path)}


def _download_failure_message(kind: str, key: str, status_code: int | None, body: str = "") -> str:
    if kind == "voice":
        base = "底层语音文件未读取成功"
        if status_code:
            base += f"（chatlog /voice 返回 HTTP {status_code}）"
        return (
            f"{base}。请先确认已安装新版程序并重新识别当前微信账号；"
            "如果仍失败，说明当前 chatlog 版本还不能解析该 WeChat 4.x 语音资源。"
        )
    if status_code:
        return f"从 chatlog 下载媒体失败: HTTP {status_code}"
    return f"从 chatlog 下载媒体失败: {body or 'unknown error'}"


async def diagnose_media(kind: str, key: str) -> dict:
    if kind not in {"voice", "image", "video"}:
        raise HTTPException(400, "不支持的媒体类型")
    if not key:
        raise HTTPException(400, "媒体 key 不能为空")

    url = _media_url(kind, key, thumb=kind in {"image", "video"})
    result = {
        "ok": False,
        "kind": kind,
        "key": key,
        "url": url,
        "chatlog_base_url": settings.chatlog_base_url,
        "chatlog_context": get_chatlog_context(),
    }

    client = shared_client()
    try:
        resp = await client.get(url, timeout=20)
        content_type = resp.headers.get("content-type", "")
        result.update(
            {
                "status_code": resp.status_code,
                "content_type": content_type,
                "content_length": len(resp.content or b""),
                "ok": resp.status_code < 400 and bool(resp.content),
            }
        )
        if resp.status_code >= 400:
            result["error"] = _download_failure_message(kind, key, resp.status_code, resp.text[:500])
            result["response_preview"] = resp.text[:500]
        elif not resp.content:
            result["error"] = "chatlog 返回了空媒体文件"
    except Exception as exc:
        result.update({"error": f"无法连接 chatlog 媒体接口: {exc}", "exception": str(exc)})

    if kind == "voice":
        result["resource_db"] = _read_voice_resource_status(key)
    return result


async def resolve_media(kind: str, key: str) -> ResolvedMedia:
    if kind not in {"voice", "image", "video"}:
        raise HTTPException(400, "不支持的媒体类型")
    if not key:
        raise HTTPException(400, "媒体 key 不能为空")

    url = _media_url(kind, key, thumb=kind in {"image", "video"})
    client = shared_client()
    try:
        resp = await client.get(url, timeout=60)
        resp.raise_for_status()
    except httpx.HTTPStatusError as exc:
        diagnostics = await diagnose_media(kind, key)
        log.warning("[media_resolver] media download failed: %s", diagnostics)
        raise HTTPException(
            502,
            {
                "message": _download_failure_message(kind, key, exc.response.status_code, exc.response.text[:500]),
                "diagnostics": diagnostics,
            },
        )
    except Exception as exc:
        diagnostics = await diagnose_media(kind, key)
        log.warning("[media_resolver] media download exception: %s", diagnostics)
        raise HTTPException(
            502,
            {
                "message": _download_failure_message(kind, key, None, str(exc)),
                "diagnostics": diagnostics,
            },
        )

    if not resp.content:
        diagnostics = await diagnose_media(kind, key)
        raise HTTPException(
            502,
            {
                "message": "chatlog 返回了空媒体文件",
                "diagnostics": diagnostics,
            },
        )

    return ResolvedMedia(
        bytes=resp.content,
        content_type=resp.headers.get("content-type", "application/octet-stream"),
        url=url,
    )