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

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(),
    }

    async with httpx.AsyncClient(timeout=20, trust_env=False, follow_redirects=True) as client:
        try:
            resp = await client.get(url)
            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"})
    async with httpx.AsyncClient(timeout=60, trust_env=False, follow_redirects=True) as client:
        try:
            resp = await client.get(url)
            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,
    )