contracts/sim/actions.py

"""
actions.py — API-обёртки для симулятора.

Все запросы к contracts.kube5s.ru с таймаутами.
Эмулирует действия пользователя через UI (загрузка, classify, groups, compare).
"""

import requests, json, time, sys, os, uuid

BASE = "https://contracts.kube5s.ru"
TIMEOUT = 120

# Убрать warnings
import urllib3
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

def _post(url, **kw):
    return requests.post(url, timeout=TIMEOUT, verify=False, **kw)

def _get(url, **kw):
    return requests.get(url, timeout=TIMEOUT, verify=False, **kw)

# ═══════════════════════════════════════════════════════════════
# Действия
# ═══════════════════════════════════════════════════════════════

class Session:
    """Состояние сессии: batch_id, contract_id, files."""
    def __init__(self):
        self.batch_id = str(uuid.uuid4())
        self.contract_id = None
        self.files = []  # {name, doc_id, status}

    def upload_file(self, path):
        """Загрузить один файл. Возвращает ответ API."""
        fname = os.path.basename(path)
        with open(path, 'rb') as f:
            data = {'batch_id': self.batch_id}
            if self.contract_id:
                data['contract_id'] = self.contract_id
            r = _post(f"{BASE}/upload", files={'files': (fname, f)}, data=data)
        result = r.json()
        if result.get('ok'):
            doc_id = result.get('doc_id')
            parsed = result.get('parsed', {})
            self.files.append({
                'name': fname, 'doc_id': doc_id,
                'status': parsed.get('status', 'uploaded'),
                'element_count': parsed.get('element_count', 0)
            })
            if result.get('contract_id'):
                self.contract_id = result['contract_id']
        return result

    def upload_files(self, paths, delay=0.2):
        """Загрузить несколько файлов последовательно."""
        for p in paths:
            r = self.upload_file(p)
            status = r.get('parsed', {}).get('status', r.get('ok', '?'))
            print(f"  ↑ {os.path.basename(p)}: {status}")
            time.sleep(delay)

    def classify(self):
        """Запустить классификацию, ждать завершения."""
        r = _post(f"{BASE}/api/classify-batch", json={'batch_id': self.batch_id})
        result = r.json()
        if not result.get('ok'):
            print(f"  ✗ classify failed: {result}")
            return result

        # Поллить прогресс
        total = result.get('total', 0)
        print(f"  ⏳ classify {total} файлов...")
        for _ in range(300):  # макс 10 минут (300 × 2с)
            time.sleep(2)
            pr = _get(f"{BASE}/api/batch-progress?batch={self.batch_id}").json()
            done = pr.get('counts', {}).get('classified', 0) + pr.get('counts', {}).get('failed', 0)
            if done >= total and total > 0:
                print(f"  ✓ classify done: {pr.get('counts')}")
                return pr
        return {"ok": False, "error": "timeout"}

    def get_groups(self):
        """Получить группы после классификации."""
        r = _get(f"{BASE}/api/groups?batch={self.batch_id}")
        data = r.json()
        groups = data.get('groups', [])
        real = [g for g in groups if g.get('contract_number') != '__unresolved__']
        unres = [g for g in groups if g.get('contract_number') == '__unresolved__']
        print(f"  📋 группы: {len(real)} (+ {len(unres)} нераспознано)")
        for g in real:
            print(f"     №{g.get('contract_number')} — {g.get('counterparty','?')} ({len(g.get('documents',[]))} док.)")
        return data

    def apply_and_compare(self, group_index=0):
        """Применить группу и запустить сравнение через SSE."""
        r = _get(f"{BASE}/api/groups?batch={self.batch_id}")
        groups = r.json().get('groups', [])
        real = [g for g in groups if g.get('contract_number') != '__unresolved__']
        if not real:
            print("  ✗ нет групп для сравнения")
            return None

        group = real[group_index]
        print(f"  ⚖ compare группы №{group.get('contract_number')}...")

        # Применить группу
        ar = _post(f"{BASE}/api/apply-groups", json={'batch_id': self.batch_id, 'groups': [group]})
        ad = ar.json()
        if not ad.get('ok') or not ad.get('contract_ids'):
            print(f"  ✗ apply failed: {ad}")
            return None

        cid = ad['contract_ids'][0]

        # SSE compare
        import sseclient  # pip install sseclient-py
        url = f"{BASE}/process-v2?contract_id={cid}"
        response = requests.get(url, stream=True, timeout=300, verify=False)
        client = sseclient.SSEClient(response)

        sections = {}
        total_ops = 0
        for event in client.events():
            d = json.loads(event.data)
            t = d.get('type')

            if t == 'extract_start':
                sections[d['supplement_id']] = d['filename']
            elif t == 'llm_done':
                print(f"     ✓ {sections.get(d['supplement_id'],'?')}: {d.get('ops_count',0)} оп. {d.get('mode','?')} ({d.get('time_s',0)}с)")
            elif t == 'applied':
                ops = d.get('ops', [])
                total_ops += len(ops)
                s = d.get('summary', {})
                print(f"     📊 +{s.get('added',0)} ~{s.get('updated',0)} -{s.get('deleted',0)}")
            elif t == 'done':
                print(f"  ✓ compare done: {d.get('total_time_s',0)}с, всего {total_ops} оп.")
                return d
            elif t == 'error':
                print(f"  ✗ compare error: {d.get('message','?')}")
                return None

        return None

    def cleanup(self):
        """Очистить БД после теста."""
        try:
            _post(f"{BASE}/api/cleanup")
        except Exception:
            pass


def upload_files_from_dir(session, dir_path, pattern="*", delay=0.2):
    """Загрузить все файлы из папки."""
    import glob
    files = sorted(glob.glob(os.path.join(dir_path, pattern)))
    files = [f for f in files if f.endswith(('.docx', '.doc', '.pdf', '.zip'))]
    if not files:
        print(f"  (нет файлов в {dir_path})")
        return
    print(f"  Загрузка {len(files)} файлов из {dir_path}...")
    session.upload_files(files, delay)