# 结构化语料数据模型
# HLDP://tools/notion-corpus-loader/structure

from dataclasses import dataclass, field
from typing import List, Dict, Optional
import json

@dataclass
class StructuredCorpusItem:
    source_id: str = ""
    source_type: str = ""
    hldp_path: str = ""
    title: str = ""
    blocks: List[Dict] = field(default_factory=list)
    raw_text: str = ""
    cognitive_jumps: List[str] = field(default_factory=list)
    causal_chains: List[str] = field(default_factory=list)
    persona: str = ""
    quality_score: float = 0.5
    quality_tags: List[str] = field(default_factory=list)
    created_time: str = ""
    last_edited_time: str = ""
    metadata: Dict = field(default_factory=dict)
    
    def to_training_sample(self, format: str = "text") -> str:
        if format == "text": return self._to_text_sample()
        elif format == "chat": return self._to_chat_sample()
        return self._to_text_sample()
    
    def _to_text_sample(self) -> str:
        parts = []
        if self.hldp_path: parts.append(f"[HLDP_PATH]{self.hldp_path}[/HLDP_PATH]")
        if self.persona: parts.append(f"[PERSONA]{self.persona}[/PERSONA]")
        if self.title: parts.append(f"[TITLE]{self.title}[/TITLE]")
        for jump in self.cognitive_jumps[:3]: parts.append(f"[COGNITIVE_JUMP]{jump}[/COGNITIVE_JUMP]")
        for chain in self.causal_chains[:3]: parts.append(f"[CAUSAL_CHAIN]{chain}[/CAUSAL_CHAIN]")
        if self.quality_score > 0.7: parts.append("[QUALITY_HIGH]")
        elif self.quality_score > 0.4: parts.append("[QUALITY_MEDIUM]")
        parts.append("[CONTENT]")
        for block in self.blocks[:50]:
            bt = block.get('type','paragraph'); c = block.get('content','')
            if c: parts.append(f"[{bt.upper()}]{c}[/{bt.upper()}]")
        parts.append("[/CONTENT]")
        return '\n'.join(parts)
    
    def _to_chat_sample(self) -> str:
        parts = []
        if self.persona: parts.append(f"[PERSONA]{self.persona}[/PERSONA]")
        parts.append("[THINKING]")
        for block in self.blocks[:30]:
            c = block.get('content','')
            if c: parts.append(c[:200])
        parts.append("[/THINKING]")
        return '\n'.join(parts)
    
    def estimate_tokens(self) -> int:
        cc = sum(1 for c in self.raw_text if '\u4e00' <= c <= '\u9fff')
        tc = len(self.raw_text)
        return int(cc * 1.5 + (tc - cc) * 0.75)

@dataclass
class CorpusCollection:
    items: List[StructuredCorpusItem] = field(default_factory=list)
    total_tokens: int = 0
    persona_distribution: Dict[str, int] = field(default_factory=dict)
    
    def add(self, item: StructuredCorpusItem):
        self.items.append(item)
        t = item.estimate_tokens()
        self.total_tokens += t
        p = item.persona or "unknown"
        self.persona_distribution[p] = self.persona_distribution.get(p,0) + t
    
    def summary(self) -> str:
        total = self.total_tokens
        s = f"{total/1e8:.2f}亿" if total > 1e8 else (f"{total/1e4:.1f}万" if total > 1e4 else str(total))
        lines = [f"条目: {len(self.items)}", f"总token: {s}", f"分布:"]
        for p, t in sorted(self.persona_distribution.items(), key=lambda x:-x[1]):
            pct = t/total*100 if total else 0
            lines.append(f"  {p}: {t/1e4:.1f}万 ({pct:.1f}%)")
        jc = sum(1 for i in self.items if i.cognitive_jumps)
        lines.append(f"含跃迁标注: {jc}/{len(self.items)}")
        return '\n'.join(lines)
    
    def to_jsonl(self, fp: str, fmt: str = "text"):
        with open(fp,'w',encoding='utf-8') as f:
            for item in self.items:
                f.write(json.dumps({"text":item.to_training_sample(fmt),"metadata":{"source_id":item.source_id,"hldp_path":item.hldp_path,"persona":item.persona,"quality_score":item.quality_score,"tokens_estimate":item.estimate_tokens()}},ensure_ascii=False)+'\n')