#!/usr/bin/env python3
"""GPU KoSimCSE 2차 침해 검토: 별도 임베딩 인덱스만 생성한다."""
from __future__ import annotations
import argparse, json, logging, os, sys, time
from collections import Counter
from pathlib import Path
import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from app.api.schemas import DetectOptions
from app.core.config import get_settings
from app.engine.detector import PlagiarismDetector
from app.engine.persistent_index import _evidence_spans, _covered_length
from app.engine.provenance import CorpusStore
from app.engine.similarity import SimilarityHit, _element_similarities
from app.engine.structural import extract_lemmas, lemma_overlap_ratio
from scripts.run_infringement_batch import build_query_plan, build_xlsx, result_row

LOG = logging.getLogger("embedding_batch")

def build_or_load(store, index_dir: Path, model_name: str, device: str, batch: int):
    """운영 hash 인덱스와 독립된 KoSimCSE 행렬을 생성한다."""
    from sentence_transformers import SentenceTransformer
    index_dir.mkdir(parents=True, exist_ok=True)
    meta_path, vec_path = index_dir / "index.json", index_dir / "embeddings.npz"
    segments = list(store.iter_segments())
    ids = [s.segment_id for s in segments]
    hashes = {s.segment_id: s.text_sha256 for s in segments}
    if meta_path.exists() and vec_path.exists():
        meta = json.loads(meta_path.read_text())
        if meta.get("segment_ids") == ids and meta.get("text_hashes") == hashes and meta.get("model") == model_name:
            return np.load(vec_path)["embeddings"], meta, SentenceTransformer(model_name, device=device)
    model = SentenceTransformer(model_name, device=device)
    LOG.info("gpu_index_encode segments=%d model=%s device=%s", len(segments), model_name, device)
    vectors = model.encode([s.text[:2048] for s in segments], batch_size=batch, normalize_embeddings=True, show_progress_bar=False, convert_to_numpy=True)
    vectors = np.asarray(vectors, dtype=np.float32)
    np.savez_compressed(vec_path, embeddings=vectors)
    meta = {"version": 1, "backend": "KoSimCSE", "model": model_name, "device": device, "segment_ids": ids, "segment_document_ids": [s.document_id for s in segments], "document_source_groups": store.document_source_groups(), "text_hashes": hashes}
    meta_path.write_text(json.dumps(meta, ensure_ascii=False), encoding="utf-8")
    return vectors, meta, model

def main():
    p=argparse.ArgumentParser(); p.add_argument("--database",type=Path,required=True); p.add_argument("--index-dir",type=Path,required=True); p.add_argument("--out-jsonl",type=Path,required=True); p.add_argument("--out-xlsx",type=Path,required=True); p.add_argument("--first-jsonl",type=Path,required=True); p.add_argument("--model",default="BM-K/KoSimCSE-roberta-multitask"); p.add_argument("--batch-size",type=int,default=64); p.add_argument("--progress-every",type=int,default=25); p.add_argument("--exclude-segments",type=Path,help="검사·색인에서 뺄 segment_id 목록 파일 (중복 제거용)"); p.add_argument("--min-exact-span",type=int,help="연속 일치 판정 기준(공백 포함 글자 수). 미지정 시 설정값"); p.add_argument("--life-writing-segments",action="store_true",help="생활수기도 자서전과 같이 세그먼트 단위로 검사"); a=p.parse_args()
    logging.basicConfig(level=logging.INFO,format="%(asctime)s %(levelname)s %(message)s")
    if os.getenv("USE_LLM_LEGAL_JUDGE", "").lower() not in ("", "0", "false", "no", "off"): raise RuntimeError("USE_LLM_LEGAL_JUDGE=false required")
    import torch
    if not torch.cuda.is_available(): raise RuntimeError("CUDA unavailable")
    device="cuda:0"; LOG.info("gpu_ready device=%s name=%s",device,torch.cuda.get_device_name(0))
    store=CorpusStore(a.database); vectors,meta,model=build_or_load(store,a.index_dir,a.model,device,a.batch_size)
    get_settings.cache_clear(); settings=get_settings().model_copy(update={"corpus_db_path":str(a.database),"persistent_index_dir":"/nonexistent","use_persistent_index":False,"use_llm_legal_judge":False})
    # 법령/태그/규칙 엔진만 재사용한다. CPU 운영 인덱스를 변경하지 않는다.
    span_min=a.min_exact_span or settings.persistent_min_exact_span
    SCORE_MIN=settings.persistent_similarity_threshold; COVERAGE_MIN=settings.persistent_min_coverage
    if a.min_exact_span:
        settings=settings.model_copy(update={"persistent_min_exact_span":a.min_exact_span})
    LOG.info("판정 기준: 유사도>=%.2f | 연속일치>=%d자 | 커버리지>=%.2f",SCORE_MIN,span_min,COVERAGE_MIN)
    detector=PlagiarismDetector(settings); groups=store.document_source_groups(); records=store.get_segments(meta["segment_ids"]); items=build_query_plan(store,life_writing_segments=a.life_writing_segments)
    # 중복 제거: 검사 대상과 색인 후보에서 동시에 뺀다. 한쪽만 빼면 지운 원고가
    # 여전히 상대로 잡혀 중복이 결과에 남는다.
    if a.exclude_segments:
        drop={x.strip() for x in a.exclude_segments.read_text(encoding="utf-8").splitlines() if x.strip()}
        before=len(items); items=[it for it in items if (it.source_segment_id or it.document_id) not in drop]
        indexed=len(meta["segment_ids"]); keep_ix=[i for i,sid in enumerate(meta["segment_ids"]) if sid not in drop]
        vectors=vectors[keep_ix]
        meta={**meta,"segment_ids":[meta["segment_ids"][i] for i in keep_ix],
              "segment_document_ids":[meta["segment_document_ids"][i] for i in keep_ix]}
        LOG.info("dedup: 질의 %d->%d, 색인 %d->%d",before,len(items),indexed,len(keep_ix))
    a.out_jsonl.parent.mkdir(parents=True,exist_ok=True); started=time.monotonic(); rows=[]
    for n,item in enumerate(items,1):
        qv=np.asarray(model.encode([item.text[:2048]],normalize_embeddings=True,show_progress_bar=False,convert_to_numpy=True)[0],dtype=np.float32); scores=vectors@qv
        allowed=np.array([d != item.document_id and (not item.source_group or groups.get(d,"") != item.source_group) for d in meta["segment_document_ids"]])
        ix=np.flatnonzero(allowed); take=ix[np.argsort(scores[ix])[-20:][::-1]]
        qe=detector._extractor.extract(item.text); ql=extract_lemmas(item.text); hits=[]; intervals=[]; doc_intervals={}
        for i in take:
            seg=records[meta["segment_ids"][int(i)]]; ev,ints,longest=_evidence_spans(item.text,seg.text); intervals.extend(ints); doc_intervals.setdefault(seg.document_id,[]).extend(ints); relem, reels=(seg.lemmas or extract_lemmas(seg.text)), (seg.elements or detector._extractor.extract(seg.text).model_dump()); es=_element_similarities(qe, type(qe)(**reels)); combined=.30*float(scores[int(i)])+.45*lemma_overlap_ratio(ql,relem)+.15*es["characters"]+.10*es["motifs"]
            hits.append((combined,seg,float(scores[int(i)]),lemma_overlap_ratio(ql,relem),es,ev,longest))
        hits.sort(key=lambda x:x[0],reverse=True); matches=[]
        # 기준값은 검색 후보를 바꾸지 않으므로, 상위 후보의 판정재료를 남겨두면
        # 어절 기준을 바꿔가며 다시 돌리지 않고 오프라인에서 재판정할 수 있다.
        candidates=[{"seg":s2.segment_id,"doc":s2.document_id,"combined":round(c,4),
                     "longest":l,"coverage":round(_covered_length(doc_intervals.get(s2.document_id,[]))/max(1,len(item.text)),4)}
                    for c,s2,_,_,_,_,l in hits[:5]]
        for combined,seg,textsim,lemmasim,es,ev,longest in hits[:5]:
            coverage=_covered_length(doc_intervals.get(seg.document_id,[]))/max(1,len(item.text));
            if combined < SCORE_MIN and longest < span_min and coverage < COVERAGE_MIN: continue
            sh=SimilarityHit(seg.segment_id,str(seg.metadata.get("document_title",seg.document_id)),combined,textsim,lemmasim,es,[])
            match=detector._to_match(sh,True,None,None).model_copy(update={"source_document_id":seg.document_id,"source_segment_id":seg.segment_id,"matched_coverage":round(coverage,4),"longest_span":longest,"evidence_spans":ev,"match_reasons":["embedding_score" if combined>=.65 else "exact_span_or_coverage"]})
            matches.append(match)
        coverage=_covered_length(intervals)/max(1,len(item.text)); legal=detector._legal_engine.assess(max_similarity=matches[0].similarity if matches else 0,coverage=coverage,longest_span=max((m.longest_span for m in matches),default=0),legal_tags=[t.tag for m in matches for t in m.tags],query_text=item.text,evidence=[])
        class R: pass
        r=R(); r.matches=matches; r.is_infringement=bool(matches); r.confidence=round(matches[0].similarity if matches else (hits[0][0] if hits else 0),4); r.legal_risk=legal; r.score_semantics=type("S",(),{"union_coverage":round(coverage,4)})(); row=result_row(item,r,retrieval_backend=f"KoSimCSE GPU ({a.model}, {device})"); row["후보 판정재료"]=candidates; rows.append(row)
        if n%a.progress_every==0: LOG.info("progress processed=%d/%d",n,len(items))
    a.out_jsonl.write_text("\n".join(json.dumps(r,ensure_ascii=False) for r in rows)+"\n",encoding="utf-8")
    build_xlsx(rows,a.out_xlsx,backend=f"KoSimCSE GPU ({a.model}, {device})")
    from openpyxl import load_workbook
    first={json.loads(x)["query_key"]:json.loads(x) for x in a.first_jsonl.read_text(encoding="utf-8").splitlines() if x.strip()}; second={r["query_key"]:r for r in rows}; f={k for k,v in first.items() if v["침해 의심"]}; s={k for k,v in second.items() if v["침해 의심"]}; wb=load_workbook(a.out_xlsx); ws=wb.create_sheet("1차_vs_2차 비교"); ws.append(["구분","건수"]); ws.append(["1차에서만",len(f-s)]); ws.append(["2차 신규",len(s-f)]); ws.append(["양쪽 공통",len(f&s)]); ws.append([]); ws.append(["2차 신규 상위 20", "가명 제목", "결합유사도", "매칭 상대 doc"])
    for r in sorted((second[k] for k in s-f),key=lambda x:x["결합유사도"],reverse=True)[:20]: ws.append([r["doc_id"],r["가명 제목"],r["결합유사도"],r["매칭 상대 doc"]])
    wb.save(a.out_xlsx); LOG.info("completed total=%d suspected=%d first_only=%d second_only=%d common=%d",len(rows),len(s),len(f-s),len(s-f),len(f&s))
if __name__=="__main__": main()