#!/usr/bin/env python3
"""
SPIKE: atom extraction + bge-m3 embeddings + Qdrant dedup.

Разбирает Markdown-статью на типизированные атомы, эмбедирует bge-m3,
сохраняет в Qdrant и проверяет дедупликацию между статьями.

bge-m3 — единственная embed-модель (dim=1024). Не переключаться на nomic.

Запуск:
    python3 spike/extract_atoms.py --article file.md [--article file2.md] --project rabota-ufa
"""

import argparse
import json
import os
import re
import sys
import time
import uuid
from dataclasses import dataclass, field

from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams, PointStruct

OLLAMA_URL = os.environ.get("OLLAMA_URL", "http://127.0.0.1:11434")
EMBED_MODEL = "bge-m3"
COLLECTION = "content_atoms"
DIM = 1024


@dataclass
class Atom:
    atom_id: str
    project: str
    type: str
    title: str
    summary: str
    body: str
    plain_text: str
    word_count: int = 0
    created_at: str = field(default_factory=lambda: time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()))

    def to_payload(self) -> dict:
        return {
            "atom_id": self.atom_id,
            "project": self.project,
            "type": self.type,
            "title": self.title,
            "summary": self.summary,
            "body": self.body,
            "plain_text": self.plain_text,
            "word_count": self.word_count,
            "created_at": self.created_at,
        }


def extract_atoms(text: str, project: str) -> list[Atom]:
    """Эвристический extract: разбивает текст на блоки, классифицирует по паттернам."""
    lines = [ln.strip() for ln in text.split("\n") if ln.strip()]
    atoms: list[Atom] = []
    seq = 0

    for ln in lines:
        if ln.startswith("#") or ln.startswith(">") or ln.startswith("http"):
            continue
        if re.match(r"^\|.*\|$", ln) or ln.startswith("---") or ln.startswith("```"):
            continue
        if len(ln) < 20:
            continue
        seq += 1
        atoms.append(Atom(
            atom_id=f"{project}_{classify(ln)}_{seq:04d}",
            project=project,
            type=classify(ln),
            title=ln[:80],
            summary=ln[:160],
            body=ln,
            plain_text=ln,
            word_count=len(re.findall(r"\S+", ln)),
        ))
    return atoms


def classify(line: str) -> str:
    low = line.lower()
    if re.search(r"\d", line) and re.search(r"[%₽$€]", line):
        return "stat"
    if "?" in line:
        return "qa"
    if re.search(r"(подпишись|подписка|перейди|заказывай|оставь|свяжись|позвони)", low):
        return "cta"
    if re.search(r"(правильно|неправильно|ошибаеш|ты не|вот почему|секрет)", low):
        return "hook"
    if re.match(r"\d+[\.\)]", line):
        return "instruction"
    if line.startswith("«") or line.startswith("«"):
        return "quote"
    if re.search(r"(vs|против|по сравнению|лучше|дороже|дешевле)", low):
        return "comparison"
    return "claim"


def embed(text: str) -> list[float]:
    """bge-m3 embed через Ollama (dim=1024)."""
    import urllib.request
    body = json.dumps({"model": EMBED_MODEL, "prompt": text}).encode()
    req = urllib.request.Request(
        f"{OLLAMA_URL}/api/embeddings", data=body,
        headers={"Content-Type": "application/json"},
    )
    with urllib.request.urlopen(req, timeout=60) as resp:
        data = json.loads(resp.read().decode())
    emb = data["embedding"]
    if len(emb) != DIM:
        raise ValueError(f"bge-m3 вернул dim={len(emb)}, ожидалось {DIM}")
    return emb


def ensure_collection(client: QdrantClient) -> None:
    if client.collection_exists(COLLECTION):
        info = client.get_collection(COLLECTION)
        print(f"  ✓ коллекция {COLLECTION} уже существует (dim={info.config.params.vectors.size})")
        return

    client.create_collection(
        collection_name=COLLECTION,
        vectors_config=VectorParams(size=DIM, distance=Distance.COSINE),
    )
    print(f"  ✓ коллекция {COLLECTION} создана (dim={DIM}, Cosine)")


def main():
    ap = argparse.ArgumentParser(description="Spike: extract + bge-m3 + Qdrant dedup")
    ap.add_argument("--article", action="append", help=".md файл (можно несколько)")
    ap.add_argument("--project", default="spike", help="project_id (домен)")
    ap.add_argument("--threshold", type=float, default=0.85, help="порог дедупа (Cosine)")
    args = ap.parse_args()

    if not args.article:
        print("Использование: --article file.md [--article file2.md ...]")
        return 1

    client = QdrantClient(host="127.0.0.1", port=6333)
    ensure_collection(client)

    all_atoms: list[Atom] = []
    for path in args.article:
        if not os.path.exists(path):
            print(f"  ✗ файл не найден: {path}")
            continue
        with open(path, encoding="utf-8") as f:
            text = f.read()
        atoms = extract_atoms(text, args.project)
        all_atoms.extend(atoms)
        print(f"  • {os.path.basename(path)}: {len(atoms)} атомов")

    if not all_atoms:
        print("Нет атомов для обработки")
        return 1

    print(f"\nЭмбеддинг bge-m3 ({len(all_atoms)} атомов)...")
    start = time.time()
    for i in range(0, len(all_atoms), 5):
        for atom in all_atoms[i:i+5]:
            emb = embed(atom.body)
            client.upsert(
                collection_name=COLLECTION,
                points=[PointStruct(
                    id=str(uuid.uuid4()),
                    vector=emb,
                    payload=atom.to_payload(),
                )],
            )
        print(f"  • {min(i+5, len(all_atoms))}/{len(all_atoms)}")
    elapsed = time.time() - start
    print(f"  ✓ загрузили {len(all_atoms)} атомов за {elapsed:.1f}с "
          f"(~{elapsed/len(all_atoms)*1000:.0f}ms/atom)")

    print(f"\nДедупликация (порог {args.threshold}):")
    dedup_hits = 0
    for atom in all_atoms:
        emb = embed(atom.body)
        res = client.query_points(
            collection_name=COLLECTION,
            query=emb,
            limit=5,
            with_payload=True,
        )
        for hit in res.points:
            if hit.payload.get("atom_id") == atom.atom_id:
                continue
            score = hit.score or 0
            if score >= args.threshold:
                dedup_hits += 1
                other = hit.payload.get("type", "?")
                print(f"  ⚠ ДУБЛЬ: {atom.type} ↔ {other} (score={score:.3f})")
                break

    print(f"\nИтог: {len(all_atoms)} атомов, {dedup_hits} дублей найдено")
    print(f"Коллекция: 127.0.0.1:6333 / {COLLECTION}")
    return 0


if __name__ == "__main__":
    sys.exit(main())