"""
HuggingFace에서 모델을 로드하여 추론하는 예제

Usage:
    from inference_example import extract_sentences
    results = extract_sentences("삼성전자의 실적이 시장 예상을 상회했다. ...")
"""

import re
from typing import List, Dict

import torch
from transformers import AutoTokenizer

from model import (
    DocumentEncoderConfig,
    DocumentEncoderForExtractiveSummarization,
    IDX_TO_ROLE,
)


def split_into_sentences(text: str) -> List[str]:
    sentences = re.split(r"(?<=[.!?])\s+", text.strip())
    return [s.strip() for s in sentences if s.strip()]


def extract_sentences(
    text: str,
    model_name_or_path: str = "./",  # 로컬 또는 HuggingFace repo ID
    top_k: int = 3,
    threshold: float = 0.5,
    device: str = None,
) -> Dict:
    """
    텍스트에서 대표문장을 추출하고 역할을 분류합니다.

    Args:
        text: 입력 텍스트 (금융 리포트 등)
        model_name_or_path: 모델 경로 또는 HuggingFace repo ID
        top_k: 추출할 최대 문장 수
        threshold: 대표문장 판단 임계값
        device: cuda 또는 cpu

    Returns:
        dict with 'sentences', 'all_scores', 'all_roles', 'selected'
    """
    device = device or ("cuda" if torch.cuda.is_available() else "cpu")

    config = DocumentEncoderConfig.from_pretrained(model_name_or_path)
    model = DocumentEncoderForExtractiveSummarization.from_pretrained(
        model_name_or_path, config=config
    )
    model = model.to(device)
    model.eval()

    tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)

    sentences = split_into_sentences(text)
    if not sentences:
        return {"sentences": [], "all_scores": [], "all_roles": [], "selected": []}

    max_sentences = config.max_sentences
    max_length = config.max_length

    padded = sentences[:max_sentences]
    num_real = len(padded)
    while len(padded) < max_sentences:
        padded.append("")

    all_input_ids, all_attention_mask = [], []
    for s in padded:
        if s:
            enc = tokenizer(s, max_length=max_length, padding="max_length", truncation=True, return_tensors="pt")
        else:
            enc = {
                "input_ids": torch.zeros(1, max_length, dtype=torch.long),
                "attention_mask": torch.zeros(1, max_length, dtype=torch.long),
            }
        all_input_ids.append(enc["input_ids"])
        all_attention_mask.append(enc["attention_mask"])

    input_ids = torch.cat(all_input_ids, dim=0).unsqueeze(0).to(device)
    attention_mask = torch.cat(all_attention_mask, dim=0).unsqueeze(0).to(device)
    document_mask = torch.zeros(1, max_sentences, device=device)
    document_mask[0, :num_real] = 1

    with torch.no_grad():
        scores, role_logits = model(input_ids, attention_mask, document_mask)

    scores_list = scores[0, :num_real].tolist()
    role_indices = role_logits[0, :num_real].argmax(dim=-1).tolist()
    roles_list = [IDX_TO_ROLE[idx] for idx in role_indices]

    selected = []
    for i, (sent, score, role) in enumerate(zip(sentences, scores_list, roles_list)):
        if score >= threshold:
            selected.append({"index": i, "sentence": sent, "score": score, "role": role})

    selected.sort(key=lambda x: x["score"], reverse=True)
    selected = selected[:top_k]
    selected.sort(key=lambda x: x["index"])

    return {
        "sentences": sentences,
        "all_scores": scores_list,
        "all_roles": roles_list,
        "selected": selected,
    }


if __name__ == "__main__":
    text = """
    삼성전자의 2024년 4분기 실적이 시장 예상을 상회했다.
    메모리 반도체 가격 상승으로 영업이익이 전분기 대비 30% 증가했다.
    HBM3E 양산이 본격화되면서 AI 반도체 시장 점유율이 확대될 전망이다.
    다만, 중국 시장의 불확실성이 여전히 리스크 요인으로 작용하고 있다.
    회사는 올해 설비 투자를 20% 확대할 계획이다.
    """

    result = extract_sentences(text, model_name_or_path="./")

    print("=" * 60)
    print("전체 문장 분석:")
    for i, (s, sc, r) in enumerate(zip(result["sentences"], result["all_scores"], result["all_roles"])):
        marker = "*" if sc >= 0.5 else " "
        print(f"  {marker} {i+1}. [{sc:.4f}] [{r:10s}] {s}")

    print(f"\n선택된 대표문장:")
    for item in result["selected"]:
        print(f"  - [{item['score']:.4f}] [{item['role']:10s}] {item['sentence']}")