#!/usr/bin/env python3
"""
批量向量化法律条文

用法:
    python scripts/vectorize_laws.py --source-id LAW-SECURITIES-2019 --batch-size 100
    python scripts/vectorize_laws.py --all-new --batch-size 100
"""

import argparse
import os
import sys
import time
from datetime import datetime

# 添加 backend 到路径
sys.path.insert(0, 'backend')

import psycopg

# 导入向量化服务
from app.services.legal_semantic_matcher import LegalSemanticMatcher

DATABASE_URL = os.environ.get('DATABASE_URL', '')
LEGAL_ARTICLE_SEMANTIC_VECTOR_VERSION = "local-token-v1"

def get_vectorization_candidates(db_url: str, source_ids: list = None, limit: int = 1000):
    """获取需要向量化的法律条文"""
    conn = psycopg.connect(db_url)
    cur = conn.cursor()

    conditions = []
    params = []

    if source_ids:
        placeholders = ','.join(['%s'] * len(source_ids))
        conditions.append(f"source_id IN ({placeholders})")
        params.extend(source_ids)

    conditions.append("(semantic_vector_json IS NULL OR semantic_vector_version != %s)")
    params.append(LEGAL_ARTICLE_SEMANTIC_VECTOR_VERSION)

    query = f"""
        SELECT id, source_id, source_title, article_no, article_text, chapter_title, section_title
        FROM legal_articles
        WHERE {' AND '.join(conditions)}
        ORDER BY id
        LIMIT %s
    """
    params.append(limit)

    cur.execute(query, params)
    rows = cur.fetchall()
    conn.close()

    return [
        {
            'id': row[0],
            'source_id': row[1],
            'source_title': row[2],
            'article_no': row[3],
            'article_text': row[4],
            'chapter_title': row[5],
            'section_title': row[6]
        }
        for row in rows
    ]

def vectorize_text(source_title: str, article_text: str, chapter_title: str = None, section_title: str = None) -> dict:
    """向量化单个条文"""
    matcher = LegalSemanticMatcher()

    # 构建向量化源文本
    source_text = '\n'.join(filter(None, [
        source_title,
        chapter_title,
        section_title,
        article_text
    ]))

    vector = matcher.vectorize(source_text)
    return vector

def update_vector(db_url: str, article_id: str, vector: dict):
    """更新向量化结果"""
    conn = psycopg.connect(db_url)
    cur = conn.cursor()

    cur.execute("""
        UPDATE legal_articles
        SET semantic_vector_json = %s,
            semantic_vector_version = %s,
            semantic_vector_updated_at = NOW()
        WHERE id = %s
    """, (json.dumps(vector, ensure_ascii=False), LEGAL_ARTICLE_SEMANTIC_VECTOR_VERSION, article_id))

    conn.commit()
    conn.close()

def main():
    parser = argparse.ArgumentParser(description='Vectorize legal articles')
    parser.add_argument('--source-id', action='append', help='Source ID(s) to vectorize')
    parser.add_argument('--all-new', action='store_true', help='Vectorize all newly imported laws')
    parser.add_argument('--batch-size', type=int, default=100, help='Batch size')
    parser.add_argument('--limit', type=int, default=5000, help='Max articles to process')
    args = parser.parse_args()

    if not DATABASE_URL:
        print("ERROR: DATABASE_URL environment variable not set")
        sys.exit(1)

    # 确定要向量化的法律
    if args.all_new:
        source_ids = [
            'LAW-LABOR-CONTRACT-IMPL-2008', 'LAW-SOCIAL-INSURANCE-2018', 'LAW-WORK-INJURY-2010',
            'LAW-PARTNERSHIP-2006', 'LAW-BANKRUPTCY-2006', 'LAW-TRUST-2001', 'LAW-SECURITIES-2019',
            'LAW-TENDER-2017', 'LAW-GOVERNMENT-PROCUREMENT-2014', 'LAW-GOVERNMENT-PROCUREMENT-IMPL-2015',
            'LAW-CONSTRUCTION-2019', 'LAW-CONSTRUCTION-QUALITY-2019', 'LAW-COPYRIGHT-2020',
            'LAW-COPYRIGHT-IMPL-2013', 'LAW-INFO-NETWORK-2013', 'LAW-CRITICAL-INFO-2021',
            'LAW-REAL-ESTATE-MGMT-2019', 'LAW-LAND-MGMT-2019', 'LAW-LAND-MGMT-IMPL-2021',
            'LAW-LABOR-DISPUTE-INTERP-2020', 'LAW-CONSTRUCTION-DISPUTE-INTERP-2020'
        ]
    elif args.source_id:
        source_ids = args.source_id
    else:
        print("ERROR: Specify --source-id or --all-new")
        sys.exit(1)

    print("=" * 60)
    print("Legal Article Vectorization")
    print("=" * 60)
    print(f"Target laws: {len(source_ids)} source IDs")
    print(f"Batch size: {args.batch_size}")
    print(f"Limit: {args.limit}")
    print("")

    start_time = time.time()
    total_processed = 0
    total_vectorized = 0
    errors = 0

    # 分批处理
    offset = 0
    while offset < args.limit:
        batch_limit = min(args.batch_size, args.limit - offset)
        candidates = get_vectorization_candidates(DATABASE_URL, source_ids, batch_limit)

        if not candidates:
            print("No more candidates to vectorize")
            break

        print(f"Processing batch: {len(candidates)} articles...")

        for article in candidates:
            total_processed += 1
            try:
                vector = vectorize_text(
                    article['source_title'],
                    article['article_text'],
                    article.get('chapter_title'),
                    article.get('section_title')
                )

                if vector:
                    update_vector(DATABASE_URL, article['id'], vector)
                    total_vectorized += 1
                    if total_vectorized % 50 == 0:
                        print(f"  Progress: {total_vectorized} vectorized...")
                else:
                    errors += 1
                    print(f"  WARNING: Empty vector for {article['source_title']} - {article['article_no']}")

            except Exception as e:
                errors += 1
                print(f"  ERROR: {article['source_title']} - {article['article_no']}: {e}")

        offset += len(candidates)

        if len(candidates) < batch_limit:
            break

    duration = time.time() - start_time

    print("")
    print("=" * 60)
    print("Vectorization Complete")
    print("=" * 60)
    print(f"Processed: {total_processed}")
    print(f"Vectorized: {total_vectorized}")
    print(f"Errors: {errors}")
    print(f"Duration: {duration:.1f} seconds")
    print(f"Speed: {total_vectorized / max(duration, 0.1):.1f} articles/second")

if __name__ == '__main__':
    import json
    main()