api.py

import os
import json
import zipfile
import json_repair
from docxtpl import DocxTemplate
from openai import OpenAI
from datetime import datetime
import chromadb
from sentence_transformers import SentenceTransformer

# --- КОНФИГУРАЦИЯ ---
API_KEY = os.getenv("DEEPSEEK_API_KEY") # Или вставь свой ключ жестко "sk-..."
BASE_URL = "https://api.deepseek.com"
TEMPLATES_DIR = "tagged_templates"
DOWNLOADS_DIR = "downloads"
REGISTRY_FILE = "templates_registry.json"
TAGS_DB_FILE = "tags_db.json"
DB_PATH = "./legal_db"
ZIP_PATH = "./legal_db.zip"
# --------------------

client = OpenAI(api_key=API_KEY, base_url=BASE_URL)

# ==============================================================================
# 1. ИНИЦИАЛИЗАЦИЯ
# ==============================================================================
collection = None
encoder = None

if not os.path.exists(DB_PATH) and os.path.exists(ZIP_PATH):
    try:
        with zipfile.ZipFile(ZIP_PATH, 'r') as z: z.extractall(".")
    except: pass

try:
    encoder = SentenceTransformer('paraphrase-multilingual-mpnet-base-v2')
    chroma_client = chromadb.PersistentClient(path=DB_PATH)
    collection = chroma_client.get_collection(name="laws")
except: pass

try:
    registry = json.load(open(REGISTRY_FILE, "r", encoding="utf-8"))
    tags_db = json.load(open(TAGS_DB_FILE, "r", encoding="utf-8"))
    clean_tags_db = {k: v for k, v in tags_db.items() if not k.startswith("_")}
except:
    registry = []
    clean_tags_db = {}

# ==============================================================================
# 2. ЛОГИКА
# ==============================================================================

async def select_best_template(user_query):
    """
    Гибридный поиск: Сначала ИИ, потом перебор ключевых слов.
    """
    # 1. Попытка через ИИ
    docs_list = ""
    for idx, item in enumerate(registry):
        # Ограничиваем длину списка, чтобы не забить контекст
        docs_list += f"{idx}. {item['filename']} ({item['description'][:100]}...)\n"

    system_prompt = f"""
Ты библиотекарь. Твоя цель - найти имя файла из списка, которое лучше всего подходит к запросу.
СПИСОК:
{docs_list}

ВЕРНИ ТОЛЬКО JSON: {{"filename": "имя_файла.docx"}}
ЕСЛИ НИЧЕГО НЕ ПОДХОДИТ, ВЕРНИ: {{"filename": null}}
"""
    try:
        response = client.chat.completions.create(
            model="deepseek-chat",
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_query}
            ],
            response_format={"type": "json_object"},
            temperature=0.0
        )
        result = json_repair.loads(response.choices[0].message.content)
        ai_choice = result.get("filename")
        
        if ai_choice:
            return ai_choice
            
    except Exception as e:
        print(f"⚠️ ИИ поиск не удался: {e}")

    # 2. ФОЛЛБЭК: Тупой поиск по словам (Если ИИ не справился)
    print("⚠️ ИИ не нашел шаблон. Запускаю поиск по ключевым словам...")
    query_words = user_query.lower().replace(",", "").split()
    # Убираем мусорные слова
    stop_words = ["хочу", "мне", "нужен", "составь", "сделай", "договор", "бланк", "образец"]
    keywords = [w for w in query_words if w not in stop_words and len(w) > 3]

    if not keywords: return None

    best_match = None
    max_hits = 0

    for item in registry:
        hits = 0
        fname = item['filename'].lower()
        desc = item.get('description', '').lower()
        
        for kw in keywords:
            # Ищем совпадения корней (очень примитивно, но работает)
            root = kw[:-2] if len(kw) > 5 else kw
            if root in fname or root in desc:
                hits += 1
        
        if hits > max_hits:
            max_hits = hits
            best_match = item['filename']

    return best_match

async def extract_data_from_chat(user_query, filename):
    system_prompt = "Извлеки данные в JSON. Даты dd.mm.yyyy.\nСХЕМА:\n"
    for key, val in clean_tags_db.items():
        system_prompt += f"- {val['tag']}: {val['description']}\n"
    
    try:
        response = client.chat.completions.create(
            model="deepseek-chat",
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_query}
            ],
            response_format={"type": "json_object"},
            temperature=0.1
        )
        return json_repair.loads(response.choices[0].message.content)
    except:
        return {}

async def consult_logic(user_text):
    # Логика консультации (если файл не найден)
    context = ""
    if collection and encoder:
        try:
            vec = encoder.encode(user_text).tolist()
            res = collection.query(query_embeddings=[vec], n_results=3)
            if res['documents'][0]:
                for txt in res['documents'][0]: context += f"{txt[:1000]}\n...\n"
        except: pass

    sys_p = "Ты юрист. Отвечай, используя законы (если есть)."
    try:
        r = client.chat.completions.create(
            model="deepseek-chat",
            messages=[{"role":"system","content":sys_p}, {"role":"user","content":f"Вопрос: {user_text}\nКонтекст: {context}"}],
            temperature=0.3
        )
        return {"type": "text", "content": r.choices[0].message.content}
    except Exception as e:
        return {"type": "text", "content": f"Ошибка: {e}"}

async def generate_doc_logic(user_text):
    # 1. Ищем шаблон
    best_filename = await select_best_template(user_text)
    
    # ЕСЛИ ШАБЛОН НЕ НАЙДЕН — ВОЗВРАЩАЕМ ТЕКСТ (КОНСУЛЬТАЦИЮ)
    if not best_filename:
        fallback_res = await consult_logic(f"Составь текст документа: {user_text}")
        # Добавляем приписку, почему так вышло
        fallback_res["content"] = "⚠️ **Шаблон файла не найден в базе.**\nЯ составил текст документа вручную:\n\n" + fallback_res["content"]
        return fallback_res

    # 2. Проверяем файл на диске
    template_path = os.path.join(TEMPLATES_DIR, best_filename)
    if not os.path.exists(template_path):
        return {"type": "text", "content": f"⚠️ Ошибка: Шаблон '{best_filename}' есть в реестре, но файл отсутствует на сервере."}

    # 3. Извлекаем данные и заполняем
    context = await extract_data_from_chat(user_text, best_filename)
    if "doc_date" not in context: context["doc_date"] = datetime.now().strftime("%d.%m.%Y")

    try:
        doc = DocxTemplate(template_path)
        doc.render(context)
        
        if not os.path.exists(DOWNLOADS_DIR): os.makedirs(DOWNLOADS_DIR)
        
        ts = datetime.now().strftime("%Y%m%d_%H%M%S")
        out_name = f"Ready_{best_filename.replace('.docx', '')[:20]}_{ts}.docx"
        out_path = os.path.join(DOWNLOADS_DIR, out_name)
        
        doc.save(out_path)
        
        return {
            "type": "file",
            "content": f"✅ Документ сформирован на основе шаблона: **{best_filename}**",
            "file_url": out_path
        }
    except Exception as e:
        return {"type": "text", "content": f"Ошибка при заполнении: {e}"}