#!/usr/bin/env python3
"""
合规审查：文章合规检查 → LLM迭代修复
从articles表读取待审文章，进行合规评分；不合格文章由LLM修复（最多3次），通过后更新选题状态为待发布
"""
import json, os, datetime, logging, sys, re
from pathlib import Path
from typing import Dict, List, Optional, Tuple
from dataclasses import dataclass, asdict

PROJECT_ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(PROJECT_ROOT))
sys.path.insert(0, str(PROJECT_ROOT / "platform" / "backend"))

from scripts.compliance_checker import check_article
try:
    from app.core.nvidia_client import call_llm
    HAVE_LLM = True
except ImportError:
    HAVE_LLM = False

from db_helper import get_topic_by_id, update_topic_status, get_active_llm_config, get_articles_by_topic, save_article
from content_cleaner import strip_thinking, strip_ai_preface, strip_thinking_html, clean_html_content
from prompt_loader import get_prompt, get_prompt_params

DATA_DIR = PROJECT_ROOT / "automation" / "data"
DRAFTS_DIR = DATA_DIR / "drafts"
LOGS_DIR = PROJECT_ROOT / "automation" / "logs"
TODAY = datetime.datetime.now().strftime("%Y-%m-%d")

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s',
                    handlers=[logging.FileHandler(LOGS_DIR / f"optimizer_{TODAY}.log"), logging.StreamHandler()])
logger = logging.getLogger(__name__)

DEFAULT_PLATFORM_TAGS = {
    "zhihu": ["科技", "职场"],
    "xiaohongshu": ["AI", "可持续", "生活方式"]
}

_cached_platform_tags = None

def _load_platform_tags():
    global _cached_platform_tags
    if _cached_platform_tags is not None:
        return _cached_platform_tags
    try:
        from app.database import SessionLocal
        from app.models import PlatformConfig
        db = SessionLocal()
        try:
            configs = db.query(PlatformConfig).filter(PlatformConfig.is_active == True).all()
            if configs:
                _cached_platform_tags = {}
                for c in configs:
                    tags = c.to_dict().get("allowed_tags", [])
                    if tags:
                        _cached_platform_tags[c.platform] = tags
                if _cached_platform_tags:
                    logger.info(f"从DB加载 {len(_cached_platform_tags)} 个平台的标签")
                    return _cached_platform_tags
        finally:
            db.close()
    except Exception as e:
        logger.warning(f"从DB加载 platform_tags 失败: {e}")
    _cached_platform_tags = DEFAULT_PLATFORM_TAGS
    return _cached_platform_tags

def get_platform_tags():
    return _load_platform_tags()

_llm_config_cache = None

def get_llm_config():
    global _llm_config_cache
    if _llm_config_cache is None:
        _llm_config_cache = get_active_llm_config()
    return _llm_config_cache

@dataclass
class OptimizationResult:
    file: str
    platform: str
    topic_id: str
    title: str
    original_issues: int
    fixed_issues: int
    final_score: int
    status: str

def load_topic_map():
    from db_helper import export_topics_to_json
    topics = export_topics_to_json()
    return {t['id']: t for t in topics}

def get_articles_from_db(topic_ids: Optional[List[str]] = None, today_only: bool = False) -> List[Tuple[str, str, str]]:
    """从 articles 表读取 HTML 内容

    Returns: [(html_content, platform, topic_id), ...]
    """
    from db_helper import get_articles_by_topic
    results = []
    seen_topics = set()
    if topic_ids:
        for tid in topic_ids:
            articles = get_articles_by_topic(tid)
            for a in articles:
                if a.get("html_content"):
                    results.append((a["html_content"], a["platform"], a["topic_id"]))
                    seen_topics.add(a["topic_id"])
    else:
        from app.database import SessionLocal
        from app.models import Article
        from sqlalchemy import func
        db = SessionLocal()
        try:
            query = db.query(Article).filter(Article.html_content.isnot(None))
            if today_only:
                cutoff = datetime.datetime.now() - datetime.timedelta(hours=24)
                query = query.filter(Article.created_at >= cutoff)
            all_articles = query.all()
            for a in all_articles:
                results.append((a.html_content, a.platform, a.topic_id))
        finally:
            db.close()
    return results

def fix_wechat_title(html: str, title: str) -> str:
    suffix = f" - {TODAY} - 微信公众号"
    max_base_len = 32 - len(suffix)
    title_tag = re.search(r'<title>([^<]+)</title>', html)
    if title_tag:
        full_title = title_tag.group(1)
        if full_title.endswith(suffix):
            base = full_title[:-len(suffix)]
        else:
            base = full_title.split(" - ")[0]
        if len(base) > max_base_len:
            base = base[:max_base_len-3] + "..."
        new_full = base + suffix
        html = html.replace(full_title, new_full)
    h1_match = re.search(r'<h1[^>]*>([^<]+)</h1>', html)
    if h1_match:
        current_h1 = h1_match.group(1)
        base_h1 = current_h1.split(" - ")[0] if " - " in current_h1 else current_h1
        if len(base_h1) > 32:
            base_h1 = base_h1[:29] + "..."
        html = html.replace(current_h1, base_h1)
    return html

def fix_tags(html: str, platform: str) -> str:
    tags_map = get_platform_tags()
    if platform == "zhihu":
        tags_str = " ".join(f"#{t}" for t in tags_map.get("zhihu", []))
        if '<div class="tags">' in html:
            old = html.split('<div class="tags">')[1].split('</div>')[0]
            html = html.replace(f'<div class="tags">{old}</div>', f'<div class="tags">{tags_str}</div>')
    elif platform == "xiaohongshu":
        tags_str = " ".join(f"#{t}" for t in tags_map.get("xiaohongshu", []))
        if '<div class="hashtags">' in html:
            old = html.split('<div class="hashtags">')[1].split('</div>')[0]
            html = html.replace(f'<div class="hashtags">{old}</div>', f'<div class="hashtags">{tags_str}</div>')
    return html

def polish_with_llm(html: str, platform: str, remaining_issues: Optional[List[Dict]] = None) -> Tuple[str, Optional[str]]:
    """用 LLM 优化文章内容，返回 (html, log_message_or_None)
    如果指定 remaining_issues，则针对性修复合规问题
    LLM 失败时自动重试一次
    LLM 提供者由 LLM_TASK_PROVIDER 环境变量决定（scheduler 从 TaskConfig 读取设置）
    """
    if not HAVE_LLM:
        return html, None
    prompt_key = "compliance_fix" if remaining_issues else "compliance_polish"
    for attempt in range(2):
        try:
            temperature = 0.5
            params = get_prompt_params(prompt_key)
            max_tokens = params.get("max_tokens", 8000)
            system_prompt = "你是一个专业的内容合规与优化助手，擅长在保持文章质量和可读性的前提下修复合规问题。"
            if remaining_issues:
                issues_desc = "\n".join(
                    f"- [{i['type']}] {i.get('category','')}: {i.get('detail','')}"
                    for i in remaining_issues
                )
                prompt = get_prompt(prompt_key, issues_desc=issues_desc, html=html)
            else:
                prompt = get_prompt(prompt_key, html=html)
            polished = call_llm(prompt, temperature=temperature, max_tokens=max_tokens, system_prompt=system_prompt)
            polished = clean_html_content(polished)
            polished = strip_ai_preface(polished)
            polished = strip_thinking_html(polished)
            if '<h2' in polished or '<p>' in polished:
                min_acceptable = min(len(html) * 0.3, 12000)
                if len(polished) > min_acceptable and len(polished) > 100:
                    tag = "针对性修复" if remaining_issues else "常规润色"
                    return polished, f"LLM {tag}"
                logger.warning(f"LLM 优化输出过短，保留原文 (len={len(polished)})")
        except Exception as e:
            logger.warning(f"LLM 优化失败 (尝试 {attempt+1}/2): {e}")
        if attempt == 0:
            logger.info(f"重试 LLM 调用...")
    return html, None

def optimize_article(html: str, platform: str, topic_data: Dict, remaining_issues: Optional[List[Dict]] = None) -> Tuple[str, List[str]]:
    logs = []
    html = strip_thinking_html(html)
    if platform == "wechat":
        html = fix_wechat_title(html, topic_data.get("title", ""))
        logs.append("标题截断（含后缀）")
    if platform in ["zhihu", "xiaohongshu"]:
        before = html
        html = fix_tags(html, platform)
        if html != before:
            tags_map = get_platform_tags()
            logs.append(f"标签标准化为{tags_map.get(platform, [])}")
    polished, pol_log = polish_with_llm(html, platform, remaining_issues)
    if pol_log:
        html = polished
        logs.append(pol_log)
    return html, logs

def _load_platform_configs() -> Dict[str, Dict]:
    """从 DB 加载所有平台配置"""
    from app.database import SessionLocal
    from app.models import PlatformConfig
    db = SessionLocal()
    try:
        configs = db.query(PlatformConfig).all()
        return {c.platform: c.to_dict() for c in configs}
    finally:
        db.close()

def main(topic_ids: List[str] = None, today_only: bool = False):
    logger.info("=== 合规审查与优化开始 ===")
    llm_provider = os.getenv("LLM_TASK_PROVIDER", "sensenova")
    logger.info(f"LLM 配置: {llm_provider} — 合规审查")

    platform_configs = _load_platform_configs()
    logger.info(f"已加载 {len(platform_configs)} 个平台配置")

    if today_only:
        logger.info("仅处理当天创建的选题文章")

    articles = get_articles_from_db(topic_ids, today_only)
    if not articles:
        logger.warning("未找到任何文章（可能尚未创作或同步到 DB）")
        report_file = DRAFTS_DIR / TODAY / "optimization_report.json"
        report_file.parent.mkdir(parents=True, exist_ok=True)
        with open(report_file, 'w', encoding='utf-8') as f:
            json.dump({
                "date": TODAY,
                "summary": {"total_articles": 0, "passed_auto": 0, "average_score": 0},
                "details": []
            }, f, ensure_ascii=False, indent=2)
        print("OPTIMIZATION_COMPLETE: 0 articles found")
        sys.exit(0)

    topic_map = load_topic_map()
    results = []
    force_passed_topics = set()

    for html, platform_dir, topic_id in articles:
        topic_data = topic_map.get(topic_id)
        if not topic_data:
            logger.warning(f"未找到选题: {topic_id}")
            continue

        pc = platform_configs.get(platform_dir, {})
        check_result = check_article(html, platform_dir, topic_data, platform_config=pc)
        issues = check_result['issues']
        score = check_result['score']
        label = f"{platform_dir}/{topic_id}"

        if issues:
            current_html = html
            current_issues = list(issues)
            current_score = score
            for attempt in range(3):
                optimized_html, opt_logs = optimize_article(current_html, platform_dir, topic_data, current_issues)
                recheck = check_article(optimized_html, platform_dir, topic_data)
                if recheck['passed'] and opt_logs and recheck['score'] > current_score:
                    save_article(topic_id, platform_dir, optimized_html)
                    logger.info(f"✅ {label} 已修复并通过审查 (第{attempt+1}次修复) 分数: {current_score}→{recheck['score']}")
                    results.append(OptimizationResult(
                        file=f"db:{platform_dir}_{topic_id}",
                        platform=platform_dir,
                        topic_id=topic_id,
                        title=topic_data.get('title',''),
                        original_issues=len(issues),
                        fixed_issues=len(issues),
                        final_score=recheck['score'],
                        status="passed"
                    ))
                    break
                if optimized_html == current_html and not opt_logs:
                    logger.warning(f"{label} LLM 未输出有效修改（第{attempt+1}次），继续重试...")
                    continue
                current_issues = recheck['issues']
                current_html = optimized_html
                current_score = recheck['score']
            else:
                save_article(topic_id, platform_dir, current_html)
                logger.warning(f"⚠️ {label} 仍有 {len(current_issues)} 个问题未修复，已强制通过（当前分 {current_score}）")
                force_passed_topics.add(topic_id)
                results.append(OptimizationResult(
                    file=f"db:{platform_dir}_{topic_id}",
                    platform=platform_dir,
                    topic_id=topic_id,
                    title=topic_data.get('title',''),
                    original_issues=len(issues),
                    fixed_issues=len(issues) - len(current_issues),
                    final_score=current_score,
                    status="force_passed"
                ))
        else:
            results.append(OptimizationResult(
                file=f"db:{platform_dir}_{topic_id}",
                platform=platform_dir,
                topic_id=topic_id,
                title=topic_data.get('title',''),
                original_issues=0,
                fixed_issues=0,
                final_score=score,
                status="passed"
            ))
            logger.info(f"✅ {label} 合规检查通过 ({score}分)")

    passed_scores = {}
    for res in results:
        if res.topic_id not in passed_scores:
            passed_scores[res.topic_id] = []
        passed_scores[res.topic_id].append(res.final_score)

    for tid, scores in passed_scores.items():
        if tid in force_passed_topics:
            logger.warning(f"选题 {tid} 仍有未修复问题，跳过状态更新（保留当前状态）")
            continue
        avg_score = sum(scores) // len(scores)
        update_topic_status(tid, 'ready', compliance_score=avg_score, reviewed_at=datetime.datetime.now())
        logger.info(f"选题 {tid} 状态 → ready（待发布）, 合规分={avg_score}")

    report_file = DRAFTS_DIR / TODAY / "optimization_report.json"
    report_file.parent.mkdir(parents=True, exist_ok=True)
    report = {
        "date": TODAY,
        "summary": {
            "total_articles": len(results),
            "passed_auto": len(results),
            "average_score": sum(r.final_score for r in results) / len(results) if results else 0
        },
        "details": [asdict(r) for r in results]
    }
    with open(report_file, 'w', encoding='utf-8') as f:
        json.dump(report, f, ensure_ascii=False, indent=2)

    logger.info(f"✅ 合规审查完成: {len(results)} 篇文章, 其中 {len([r for r in results if r.status=='force_passed'])} 篇强制通过")
    print(f"OPTIMIZATION_COMPLETE: {len(results)} articles, {len(force_passed_topics)} topics skipped status update")
    sys.exit(0)

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description='合规审查与优化任务')
    parser.add_argument('--topic-ids', help='逗号分隔的选题ID列表，例如: A01,B02')
    parser.add_argument('--today-only', action='store_true', help='仅处理当天创建的选题文章')
    args = parser.parse_args()
    topic_ids = args.topic_ids.split(',') if args.topic_ids else None
    main(topic_ids, today_only=args.today_only)