178 lines
8.2 KiB
Python
178 lines
8.2 KiB
Python
import re
|
||
from typing import Optional
|
||
from jinja2 import Environment, FileSystemLoader
|
||
|
||
from pt_gen.core.config import Settings
|
||
from pt_gen.models.movie import MovieInfo
|
||
from pt_gen.services.cache import RedisCache
|
||
from pt_gen.services.douban import DoubanScraper
|
||
from pt_gen.services.tmdb import TMDBClient
|
||
from pt_gen.services.uploader import ImageUploader
|
||
|
||
class InfoOrchestrator:
|
||
def __init__(self, settings: Settings):
|
||
self.settings = settings
|
||
self.tmdb = TMDBClient(api_key=self.settings.tmdb.api_key)
|
||
self.douban = DoubanScraper(cookie=self.settings.douban.cookie)
|
||
self.cache = RedisCache(
|
||
host=self.settings.redis.host,
|
||
port=self.settings.redis.port,
|
||
db=self.settings.redis.db
|
||
)
|
||
if self.settings.uploader.enable:
|
||
self.uploader = ImageUploader(
|
||
api_url=self.settings.uploader.api_url,
|
||
api_key=self.settings.uploader.api_key
|
||
)
|
||
else:
|
||
self.uploader = None
|
||
|
||
self.jinja_env = Environment(loader=FileSystemLoader('templates'))
|
||
|
||
def _extract_douban_id(self, url: str) -> Optional[str]:
|
||
match = re.search(r'douban\.com/subject/(\d+)', url)
|
||
return match.group(1) if match else None
|
||
|
||
def _extract_imdb_id(self, url: str) -> Optional[str]:
|
||
match = re.search(r'imdb\.com/title/(tt\d+)', url)
|
||
return match.group(1) if match else None
|
||
|
||
def _extract_tmdb_id(self, url: str) -> Optional[str]:
|
||
match = re.search(r'themoviedb\.org/(?:movie|tv)/(\d+)', url)
|
||
return match.group(1) if match else None
|
||
|
||
async def generate_info(self, url: str) -> Optional[str]:
|
||
# --- 步骤 1: 识别URL类型并提取初始ID ---
|
||
douban_id = self._extract_douban_id(url)
|
||
imdb_id = self._extract_imdb_id(url)
|
||
tmdb_id = self._extract_tmdb_id(url)
|
||
|
||
if not (douban_id or imdb_id or tmdb_id):
|
||
return "无效或不支持的链接。请输入有效的豆瓣、IMDb或TMDB链接。"
|
||
|
||
# --- 步骤 2: ID 补全 (以 IMDb ID 为核心) ---
|
||
# 如果从 TMDB 链接开始，先找到 IMDb ID
|
||
if tmdb_id and not imdb_id:
|
||
details = await self.tmdb.get_movie_details(tmdb_id)
|
||
if details and details.get('imdb_id'):
|
||
imdb_id = details['imdb_id']
|
||
|
||
# 如果从豆瓣链接开始，爬取豆瓣页面来找到 IMDb ID
|
||
if douban_id and not imdb_id:
|
||
# 临时爬一次，只为获取 IMDb ID
|
||
temp_douban_data = await self.douban.scrape_movie_info(douban_id)
|
||
if temp_douban_data.get('imdb_id'):
|
||
imdb_id = temp_douban_data['imdb_id']
|
||
|
||
# --- 步骤 3: 查缓存 (使用唯一的 IMDb ID 作为 key) ---
|
||
if imdb_id:
|
||
cache_key = f"movieinfo:{imdb_id}"
|
||
cached_result = await self.cache.get(cache_key)
|
||
if cached_result:
|
||
print(f"命中缓存: {cache_key}")
|
||
return cached_result
|
||
else:
|
||
# 如果没有IMDb ID (例如一个没有IMDB链接的冷门豆瓣电影), 则用豆瓣ID做缓存key
|
||
cache_key = f"movieinfo:douban:{douban_id}"
|
||
|
||
|
||
# --- 步骤 4: 分别获取数据 ---
|
||
douban_data = {}
|
||
tmdb_data = {}
|
||
|
||
# 如果有豆瓣ID，就从豆瓣获取数据
|
||
if douban_id:
|
||
douban_data = await self.douban.scrape_movie_info(douban_id)
|
||
if not douban_data:
|
||
return "从豆瓣获取信息失败，请检查链接或Cookie。"
|
||
|
||
# 如果有IMDb ID，就从TMDB获取数据
|
||
if imdb_id:
|
||
tmdb_find_data = await self.tmdb.find_by_imdb_id(imdb_id)
|
||
if tmdb_find_data and tmdb_find_data.get('movie_results'):
|
||
# 确保 tmdb_id 被正确设置
|
||
tmdb_id = tmdb_find_data['movie_results'][0]['id']
|
||
tmdb_data = await self.tmdb.get_movie_details(tmdb_id)
|
||
if imdb_id and not tmdb_id:
|
||
tmdb_find_data = await self.tmdb.find_by_imdb_id(imdb_id)
|
||
if tmdb_find_data and tmdb_find_data.get('movie_results'):
|
||
tmdb_id = tmdb_find_data['movie_results'][0]['id']
|
||
|
||
if tmdb_id:
|
||
# 现在 get_movie_details 返回的是干净的数据
|
||
tmdb_data = await self.tmdb.get_movie_details(tmdb_id)
|
||
|
||
# --- 步骤 5: 智能合并数据到 MovieInfo 对象 ---
|
||
movie = MovieInfo()
|
||
# 优先使用TMDB数据，因为通常更规范
|
||
if tmdb_data:
|
||
movie.tmdb_id = tmdb_data.get('id')
|
||
movie.imdb_id = tmdb_data.get('imdb_id')
|
||
movie.original_title = tmdb_data.get('original_title')
|
||
movie.chinese_title = tmdb_data.get('title')
|
||
movie.aka_titles = tmdb_data.get('aka_titles', [])
|
||
movie.year = tmdb_data.get('release_date', '')[:4]
|
||
movie.release_date = tmdb_data.get('release_date')
|
||
movie.runtime = tmdb_data.get('runtime')
|
||
movie.spoken_languages = tmdb_data.get('spoken_languages', [])
|
||
movie.tagline = tmdb_data.get('tagline')
|
||
movie.imdb_rating = f"{tmdb_data.get('vote_average', 0):.1f}/10"
|
||
movie.directors = tmdb_data.get('directors', [])
|
||
movie.writers = tmdb_data.get('writers', [])
|
||
movie.actors = tmdb_data.get('actors', [])
|
||
movie.genres = tmdb_data.get('genres', [])
|
||
movie.countries = tmdb_data.get('countries', [])
|
||
# TMDB的简介可能是空的，先用着，后面会被豆瓣覆盖
|
||
movie.synopsis = tmdb_data.get('overview', '')
|
||
if tmdb_data.get('poster_path'):
|
||
movie.poster_url = f"https://image.tmdb.org/t/p/original{tmdb_data['poster_path']}"
|
||
|
||
|
||
# 使用豆瓣数据进行补充或覆盖
|
||
if douban_data:
|
||
movie.douban_id = douban_id
|
||
movie.douban_link = f"https://movie.douban.com/subject/{douban_id}/"
|
||
# 智能填充：如果TMDB没值，才用豆瓣的
|
||
movie.chinese_title = movie.chinese_title or douban_data.get('chinese_title')
|
||
movie.year = movie.year or douban_data.get('year')
|
||
movie.runtime = movie.runtime or douban_data.get('runtime')
|
||
movie.spoken_languages = movie.spoken_languages or douban_data.get('spoken_languages', [])
|
||
movie.directors = movie.directors or douban_data.get('directors', [])
|
||
movie.writers = movie.writers or douban_data.get('writers', [])
|
||
movie.actors = movie.actors or douban_data.get('actors', [])
|
||
movie.genres = movie.genres or douban_data.get('genres', [])
|
||
movie.countries = movie.countries or douban_data.get('countries', [])
|
||
movie.poster_url = movie.poster_url or douban_data.get('poster_url')
|
||
# 强制覆盖：豆瓣的简介、评分和获奖通常更符合中文区需求
|
||
if douban_data.get('synopsis'):
|
||
movie.synopsis = douban_data.get('synopsis')
|
||
movie.douban_rating = douban_data.get('douban_rating', 'N/A')
|
||
movie.awards = douban_data.get('awards', [])
|
||
# 确保IMDb链接存在
|
||
if movie.imdb_id:
|
||
movie.imdb_link = f"https://www.imdb.com/title/{movie.imdb_id}/"
|
||
|
||
# 确保标题有备用值
|
||
if not movie.original_title:
|
||
movie.original_title = movie.chinese_title
|
||
if not movie.chinese_title:
|
||
movie.chinese_title = movie.original_title
|
||
|
||
# --- 步骤 6: 后续处理 (图片转存、渲染、缓存) ---
|
||
if self.uploader and movie.poster_url:
|
||
print(f"开始转存图片: {movie.poster_url}")
|
||
movie.poster_url = await self.uploader.upload(movie.poster_url)
|
||
print(f"转存成功，新URL: {movie.poster_url}")
|
||
|
||
template = self.jinja_env.get_template('description.jinja2')
|
||
final_text = template.render(movie.dict())
|
||
|
||
# 写入缓存
|
||
await self.cache.set(cache_key, final_text, ttl=self.settings.redis.cache_ttl_seconds)
|
||
|
||
print("\n" + "="*40 + " FINAL RENDERED OUTPUT " + "="*40)
|
||
print(final_text)
|
||
print("="*103 + "\n")
|
||
|
||
return final_text
|