#!/usr/bin/env python3
"""
Интеллектуальное юридическое ядро МКД «Домовой 2.0» (AI Legal Engine)
Работает на базе Gemini 2.5 Flash / Pro:
- Нативный слух (Gemini Multimodal Audio) без Whisper
- Прямой асинхронный REST API (высокая отказоустойчивость без конфликтов pydantic)
- Разграничение ролей: resident (соседский помощник), lawyer (претензионно-судебный), admin (глубокий аудит)
"""

import os
import sys
import logging
import asyncio
import aiohttp
import base64
from datetime import datetime
from pathlib import Path

log = logging.getLogger("ai_legal_engine")
class AILegalEngine:
    def __init__(self):
        raw_key = os.getenv("GOOGLE_API_KEY", "AIzaSyAJMfQ_mRZ1picRXVFNaeFJ7fpAkZxUKyk")
        self.api_key = raw_key.strip("\"'") if raw_key else ""
        self.client = None
        self.model_name = "gemini-2.5-flash"
        self.deep_model_name = "gemini-2.5-flash"
        self._init_sdk_client()

    def _init_sdk_client(self):
        if not self.api_key:
            log.warning("Gemini API Key missing! AI Legal Engine will run in offline mode.")
            return

        try:
            from google import genai
            self.client = genai.Client(api_key=self.api_key)
            log.info("Gemini SDK Client initialized.")
        except Exception as e:
            log.info(f"Using high-performance direct REST engine (SDK not active: {e})")

    def get_role_instructions(self, role: str) -> str:
        if role == "resident":
            return (
                "РОЛЬ ДЛЯ ОТВЕТА: 'Соседский помощник для жильца'.\n"
                "СТИЛЬ: Спокойный, уважительный, без сложного бюрократического языка и канцелярита.\n"
                "ТРЕБОВАНИЯ:\n"
                "1. Объясни ситуацию простыми словами за 2-3 коротких абзаца.\n"
                "2. Дай пошаговый житейский совет (Шаг 1, Шаг 2, Шаг 3).\n"
                "3. Укажи, кто несет ответственность (УК, застройщик, водоканал, совет дома).\n"
                "4. Если нужно подать заявление — приведи краткий текст (3-4 строки), который жилец может скопировать и отправить в УК."
            )
        elif role == "lawyer":
            return (
                "РОЛЬ ДЛЯ ОТВЕТА: 'Профессиональный юрист по жилищному праву и защите прав потребителей'.\n"
                "СТИЛЬ: Строгий, юридически выверенный, аргументированный.\n"
                "ТРЕБОВАНИЯ:\n"
                "1. Четкая правовая квалификация действий/бездействия лиц.\n"
                "2. Прямые ссылки на статьи законов (ЖК РФ, ГК РФ, 59-ФЗ, ПП РФ № 354, ПП РФ № 491, Госстрой № 170, СанПиН 2.1.3684-21).\n"
                "3. Оценка перспектив в суде или надзорных органах (ГЖИ, Прокуратура, Роспотребнадзор).\n"
                "4. Готовый проект претензии / жалобы с правовой мотивировкой."
            )
        elif role == "admin":
            return (
                "РОЛЬ ДЛЯ ОТВЕТА: 'Главный аналитик Совета Дома и стратег судебных процессов'.\n"
                "СТИЛЬ: Стратегический, системный, аналитический.\n"
                "ТРЕБОВАНИЯ:\n"
                "1. Полный расклад с точки зрения системной выгоды для всего дома.\n"
                "2. Выявление уловок УК, скрытых манипуляций со сметами, протоколами ОСС или объемами коммунальных ресурсов.\n"
                "3. Сравнение с прецедентами и судебной практикой Ульяновского областного суда и Верховного Суда РФ.\n"
                "4. Стратегия принуждения госорганов к исполнению служебных обязанностей по 59-ФЗ и Приказу ГП № 45."
            )
        return "Отвечай вежливо и по делу."

    async def _query_rest_api(self, prompt: str, system_instruction: str, model: str, audio_path: str = None) -> str:
        """Прямой асинхронный REST запрос к Google Gemini API."""
        url = f"https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent?key={self.api_key}"
        parts = []

        # Аудио вложение
        if audio_path and os.path.exists(audio_path):
            with open(audio_path, "rb") as af:
                audio_b64 = base64.b64encode(af.read()).decode("utf-8")
            parts.append({
                "inlineData": {
                    "mimeType": "audio/ogg",
                    "data": audio_b64
                }
            })
            if not prompt:
                prompt = "Внимательно прослушай аудиосообщение жильца дома. Точно выдели юридическую суть и дай обоснованный ответ."

        parts.append({"text": prompt})

        payload = {
            "system_instruction": {
                "parts": [{"text": system_instruction}]
            },
            "contents": [
                {"parts": parts}
            ],
            "generationConfig": {
                "temperature": 0.2 if model == self.deep_model_name else 0.4
            }
        }

        async with aiohttp.ClientSession() as session:
            async with session.post(url, json=payload, timeout=aiohttp.ClientTimeout(total=45)) as resp:
                if resp.status != 200:
                    err_text = await resp.text()
                    raise RuntimeError(f"Gemini API HTTP {resp.status}: {err_text[:300]}")
                data = await resp.json()
                return data["candidates"][0]["content"]["parts"][0]["text"]

    async def query_notebooklm(self, prompt: str) -> dict:
        """Прямой запрос к блокноту NotebookLM «Разбор ответов строго по ЖК РФ»."""
        import shutil
        notebook_id = os.getenv("NOTEBOOKLM_ID", "00ed7d92-1ab2-4159-ac72-1d0c4be21377")
        candidates = [
            "/usr/local/bin/notebooklm",
            shutil.which("notebooklm"),
            "/home/matrixhasyou/swarm-services/domovoy-bot/venv-ai/bin/notebooklm",
            "notebooklm"
        ]
        nb_bin = None
        for b in candidates:
            if b and os.path.exists(b) and os.access(b, os.X_OK):
                # Проверим, не сломан ли шебанг
                try:
                    with open(b, "rb") as f:
                        first_line = f.readline()
                        if first_line.startswith(b"#!"):
                            interp = first_line[2:].strip().decode("utf-8", errors="ignore")
                            if not os.path.exists(interp):
                                continue
                except Exception:
                    pass
                nb_bin = b
                break
        if not nb_bin:
            nb_bin = "/usr/local/bin/notebooklm" if os.path.exists("/usr/local/bin/notebooklm") else "notebooklm"

        cmd = [
            nb_bin, "ask",
            "-n", notebook_id,
            prompt
        ]
        env = os.environ.copy()
        env["HOME"] = "/root" if os.path.exists("/root/.notebooklm") else os.environ.get("HOME", "/root")
        env["SSL_CERT_FILE"] = "/etc/ssl/certs/ca-certificates.crt"

        try:
            proc = await asyncio.create_subprocess_exec(
                *cmd,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
                env=env
            )
            try:
                stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=90.0)
            except asyncio.TimeoutError:
                try: proc.kill()
                except: pass
                return {"success": False, "error": "Таймаут обращения к NotebookLM (более 90 сек)", "source": "notebooklm"}

            out_str = stdout.decode('utf-8', errors='ignore').strip()
            err_str = stderr.decode('utf-8', errors='ignore').strip()

            if proc.returncode == 0 and out_str and "Authentication Error" not in out_str and "Not logged in" not in out_str:
                lines = [l for l in out_str.splitlines() if not l.startswith("Continuing conversation") and not l.startswith("Resumed conversation")]
                cleaned = "\n".join(lines).strip()
                if cleaned.startswith("Answer:"):
                    cleaned = cleaned[len("Answer:"):].strip()
                return {"success": True, "text": cleaned, "source": "notebooklm"}
            else:
                log.warning(f"NotebookLM non-zero or auth err: code={proc.returncode}, out={out_str[:200]}, err={err_str[:200]}")
                return {"success": False, "error": err_str or out_str, "source": "notebooklm"}
        except Exception as e:
            log.error(f"NotebookLM exec exception: {e}")
            return {"success": False, "error": str(e), "source": "notebooklm"}

    async def query(self, prompt: str, role: str = "resident", audio_path: str = None) -> dict:
        """Основной метод запроса: NotebookLM (первый безусловный приоритет) -> Gemini (fallback/аудио)."""
        start_time = datetime.now()

        # 1. Текстовые вопросы ВСЕХ пользователей идут строго в блокнот «Разбор ответов строго по ЖК РФ»!
        if not audio_path and prompt:
            nlm_res = await self.query_notebooklm(prompt)
            if nlm_res.get("success"):
                latency = (datetime.now() - start_time).total_seconds()
                return {
                    "success": True,
                    "text": nlm_res["text"],
                    "answer": nlm_res["text"],
                    "role": role,
                    "latency_sec": round(latency, 2),
                    "model": "NotebookLM: Разбор ответов строго по ЖК РФ"
                }
            else:
                log.warning(f"NotebookLM query failed: {nlm_res.get('error')}. Falling back to Gemini.")

        # 2. Резервный контур через Gemini API (для аудиосообщений либо если блокнот временно недоступен)
        if not self.api_key:
            return {
                "success": False,
                "text": "ИИ-модуль временно недоступен. Проверьте подключение.",
                "role": role,
                "latency_sec": 0.0
            }

        system_instruction = (
            "Ты — интеллектуальный юридический щит 10-подъездного многоквартирного дома по адресу: г. Ульяновск, пр-кт Ленинского Комсомола, д. 37 (МКД ЛенКом 37).\n"
            "Дом ведет системную юридическую защиту прав собственников против нарушений управляющей компании, бездействия надзорных органов и фальсификаций.\n"
            f"{self.get_role_instructions(role)}"
        )

        model_to_use = self.model_name

        try:
            reply_text = await self._query_rest_api(
                prompt=prompt,
                system_instruction=system_instruction,
                model=model_to_use,
                audio_path=audio_path
            )

            latency = (datetime.now() - start_time).total_seconds()
            
            # Если для админа не сработал блокнот, честно добавляем статус
            if role == "admin" and not audio_path:
                prefix = "📚 <b>[Источник: База законов РФ / Резервный шлюз Gemini]</b>\n<i>(Сессия NotebookLM требует обновления куков)</i>\n\n"
                reply_text = prefix + reply_text

            return {
                "success": True,
                "text": reply_text.strip(),
                "answer": reply_text.strip(),
                "role": role,
                "latency_sec": round(latency, 2),
                "model": model_to_use
            }

        except Exception as e:
            latency = (datetime.now() - start_time).total_seconds()
            log.error(f"Gemini Error: {e}")
            return {
                "success": False,
                "text": f"Ошибка обработки запроса ИИ: {e}",
                "role": role,
                "latency_sec": round(latency, 2)
            }

# Глобальный синглтон юридического ядра для роутеров бота
legal_engine = AILegalEngine()

if __name__ == "__main__":
    async def test():
        engine = AILegalEngine()
        res = await engine.query("Какие сроки ответа прокуратуры по закону 59-ФЗ?", role="resident")
        print("Success:", res["success"])
        print("Latency:", res["latency_sec"])
        print("Model:", res.get("model"))
        print("Text snippet:\n", res["text"][:300])

    asyncio.run(test())