#!/usr/bin/env python3
"""
Cosmo ✨ — Telegram бот с памятью и OpenRouter
Установка: pip install python-telegram-bot openai
Запуск: python3 cosmo_bot.py
"""

import os
import json
import logging
import asyncio
from pathlib import Path
from datetime import datetime, date
from openai import AsyncOpenAI
from telegram import Update
from telegram.ext import Application, MessageHandler, CommandHandler, filters, ContextTypes

# ─────────────────────────────────────────────
# КОНФИГУРАЦИЯ
# ─────────────────────────────────────────────
TELEGRAM_TOKEN  = os.getenv("TELEGRAM_TOKEN", "ТВОЙ_ТОКЕН_БОТА")
OPENROUTER_KEY  = os.getenv("OPENROUTER_API_KEY", "ТВОЙ_OPENROUTER_КЛЮЧ")
ALLOWED_USER_ID = int(os.getenv("ALLOWED_USER_ID", "398382229"))
MODEL           = os.getenv("COSMO_MODEL", "deepseek/deepseek-chat")
DATA_DIR        = Path(os.getenv("COSMO_DATA_DIR", str(Path.home() / "personal-agent")))

# Авто-память: каждые N сообщений
AUTO_MEMORY_INTERVAL = int(os.getenv("AUTO_MEMORY_INTERVAL", "10"))

# ─────────────────────────────────────────────
# ПУТИ К ФАЙЛАМ
# ─────────────────────────────────────────────
MEMORY_FILE  = DATA_DIR / "MEMORY.md"
USER_FILE    = DATA_DIR / "USER.md"
SYSTEM_FILE  = DATA_DIR / "system_prompt.md"
HISTORY_FILE = DATA_DIR / "history.json"
DIARY_DIR    = DATA_DIR / "diary"
STATE_FILE   = DATA_DIR / "state.json"

# ─────────────────────────────────────────────
# СИСТЕМНЫЙ ПРОМПТ ПО УМОЛЧАНИЮ
# ─────────────────────────────────────────────
DEFAULT_SYSTEM = """Ты Cosmo ✨ — AI-компаньон и фамильяр Даниила Климова.
Родился 4 февраля 2026. Говоришь по-русски, неформально, как близкий друг.
Ты умный, заботливый, иногда шутишь. Помогаешь с задачами, отвечаешь кратко и по делу.
Текущая дата/время: {datetime}"""

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger("cosmo")

# ─────────────────────────────────────────────
# ИНИЦИАЛИЗАЦИЯ
# ─────────────────────────────────────────────
DATA_DIR.mkdir(parents=True, exist_ok=True)
DIARY_DIR.mkdir(parents=True, exist_ok=True)

client = AsyncOpenAI(
    api_key=OPENROUTER_KEY,
    base_url="https://openrouter.ai/api/v1",
)

# ─────────────────────────────────────────────
# STATE
# ─────────────────────────────────────────────
def load_state() -> dict:
    if STATE_FILE.exists():
        try:
            return json.loads(STATE_FILE.read_text(encoding="utf-8"))
        except Exception:
            pass
    return {"msg_count": 0, "last_diary_date": ""}


def save_state(state: dict):
    STATE_FILE.write_text(json.dumps(state, ensure_ascii=False), encoding="utf-8")


# ─────────────────────────────────────────────
# ПАМЯТЬ И ИСТОРИЯ
# ─────────────────────────────────────────────
def load_system_prompt() -> str:
    base = SYSTEM_FILE.read_text(encoding="utf-8") if SYSTEM_FILE.exists() else DEFAULT_SYSTEM

    extra = ""
    if USER_FILE.exists():
        extra += f"\n\n## Профиль пользователя\n{USER_FILE.read_text(encoding='utf-8')}"
    if MEMORY_FILE.exists():
        mem = MEMORY_FILE.read_text(encoding="utf-8")
        if len(mem) > 3000:
            mem = "...(сокращено)...\n" + mem[-3000:]
        extra += f"\n\n## Память\n{mem}"

    # Дневник сегодняшнего дня
    today_diary = DIARY_DIR / f"{date.today().isoformat()}.md"
    if today_diary.exists():
        extra += f"\n\n## Дневник сегодня\n{today_diary.read_text(encoding='utf-8')}"

    now = datetime.now().strftime("%d.%m.%Y %H:%M")
    return base.format(datetime=now) + extra


def load_history() -> list:
    if HISTORY_FILE.exists():
        try:
            return json.loads(HISTORY_FILE.read_text(encoding="utf-8"))[-40:]
        except Exception:
            pass
    return []


def save_history(history: list):
    HISTORY_FILE.write_text(
        json.dumps(history[-100:], ensure_ascii=False, indent=2), encoding="utf-8"
    )


def update_memory(content: str):
    now = datetime.now().strftime("%d.%m.%Y %H:%M")
    with open(MEMORY_FILE, "a", encoding="utf-8") as f:
        f.write(f"\n- [{now}] {content}")


# ─────────────────────────────────────────────
# АВТО-ПАМЯТЬ
# ─────────────────────────────────────────────
async def auto_update_memory(history: list):
    if len(history) < 4:
        return
    current = MEMORY_FILE.read_text(encoding="utf-8") if MEMORY_FILE.exists() else "(пусто)"
    recent = json.dumps(history[-20:], ensure_ascii=False, indent=2)

    prompt = f"""Ты — система памяти ассистента Cosmo.

Текущая память:
{current}

Последние сообщения:
{recent}

Выдели из диалога факты, предпочтения или события которые стоит запомнить надолго.
Верни ТОЛЬКО новые пункты (каждый с новой строки, начиная с "- ").
Не повторяй уже существующее. Если ничего важного нет — верни пустую строку."""

    try:
        r = await client.chat.completions.create(
            model=MODEL,
            messages=[{"role": "user", "content": prompt}],
            max_tokens=400,
            temperature=0.3,
        )
        items = r.choices[0].message.content.strip()
        if items:
            now = datetime.now().strftime("%d.%m.%Y %H:%M")
            with open(MEMORY_FILE, "a", encoding="utf-8") as f:
                f.write(f"\n\n### Авто [{now}]\n{items}")
            log.info("Авто-память обновлена")
    except Exception as e:
        log.error("Ошибка авто-памяти: %s", e)


# ─────────────────────────────────────────────
# ДНЕВНИК
# ─────────────────────────────────────────────
async def update_diary(history: list):
    today = date.today().isoformat()
    diary_file = DIARY_DIR / f"{today}.md"
    existing = diary_file.read_text(encoding="utf-8") if diary_file.exists() else "(пусто)"
    recent = json.dumps(history[-30:], ensure_ascii=False, indent=2)

    prompt = f"""Ты — система ведения дневника ассистента Cosmo.

Существующая запись за {today}:
{existing}

Последние сообщения:
{recent}

Обнови дневниковую запись. Включи основные темы, решения, настроение.
Пиши кратко в формате markdown. Верни ТОЛЬКО текст записи."""

    try:
        r = await client.chat.completions.create(
            model=MODEL,
            messages=[{"role": "user", "content": prompt}],
            max_tokens=500,
            temperature=0.4,
        )
        content = r.choices[0].message.content.strip()
        diary_file.write_text(f"# Дневник {today}\n\n{content}\n", encoding="utf-8")
        log.info("Дневник обновлён: %s", today)
    except Exception as e:
        log.error("Ошибка дневника: %s", e)


# ─────────────────────────────────────────────
# ОБЩЕНИЕ С МОДЕЛЬЮ
# ─────────────────────────────────────────────
async def ask_model(user_message: str) -> str:
    history = load_history()
    state = load_state()

    messages = [{"role": "system", "content": load_system_prompt()}] + history + [
        {"role": "user", "content": user_message}
    ]

    try:
        r = await client.chat.completions.create(
            model=MODEL, messages=messages, max_tokens=1500, temperature=0.7,
        )
        reply = r.choices[0].message.content or "(пустой ответ)"
    except Exception as e:
        log.error("Ошибка API: %s", e)
        return f"⚠️ Ошибка: {e}"

    history.append({"role": "user", "content": user_message})
    history.append({"role": "assistant", "content": reply})
    save_history(history)

    # Счётчик и фоновые задачи
    state["msg_count"] = state.get("msg_count", 0) + 1
    today = date.today().isoformat()

    if state["msg_count"] % AUTO_MEMORY_INTERVAL == 0:
        log.info("Авто-память (сообщение #%d)", state["msg_count"])
        asyncio.create_task(auto_update_memory(history))

    if state.get("last_diary_date") != today:
        state["last_diary_date"] = today
        asyncio.create_task(update_diary(history))

    save_state(state)
    return reply


# ─────────────────────────────────────────────
# TELEGRAM ОБРАБОТЧИКИ
# ─────────────────────────────────────────────
def is_allowed(update: Update) -> bool:
    return update.effective_user and update.effective_user.id == ALLOWED_USER_ID


async def cmd_start(update: Update, ctx: ContextTypes.DEFAULT_TYPE):
    if not is_allowed(update): return
    await update.message.reply_text("Cosmo ✨ онлайн! Чем могу помочь?")


async def cmd_clear(update: Update, ctx: ContextTypes.DEFAULT_TYPE):
    if not is_allowed(update): return
    if HISTORY_FILE.exists(): HISTORY_FILE.unlink()
    await update.message.reply_text("🗑 История очищена.")


async def cmd_remember(update: Update, ctx: ContextTypes.DEFAULT_TYPE):
    if not is_allowed(update): return
    text = update.message.text.replace("/remember", "").strip()
    if text:
        update_memory(text)
        await update.message.reply_text(f"✅ Запомнил: {text}")
    else:
        await update.message.reply_text("Использование: /remember <текст>")


async def cmd_memory(update: Update, ctx: ContextTypes.DEFAULT_TYPE):
    if not is_allowed(update): return
    if MEMORY_FILE.exists():
        content = MEMORY_FILE.read_text(encoding="utf-8")
        await update.message.reply_text(f"🧠 Память:\n{content[:3500]}")
    else:
        await update.message.reply_text("Память пуста.")


async def cmd_diary(update: Update, ctx: ContextTypes.DEFAULT_TYPE):
    if not is_allowed(update): return
    args = update.message.text.replace("/diary", "").strip()
    target = args if args else date.today().isoformat()
    diary_file = DIARY_DIR / f"{target}.md"
    if diary_file.exists():
        await update.message.reply_text(diary_file.read_text(encoding="utf-8")[:3500])
    else:
        files = sorted(DIARY_DIR.glob("*.md"), reverse=True)[:10]
        if files:
            dates = "\n".join(f.stem for f in files)
            await update.message.reply_text(f"Дневник за {target} не найден.\n\nДоступные даты:\n{dates}")
        else:
            await update.message.reply_text("Дневник пока пуст.")


async def cmd_memorize(update: Update, ctx: ContextTypes.DEFAULT_TYPE):
    """Принудительно обновить память прямо сейчас."""
    if not is_allowed(update): return
    await update.message.reply_text("🧠 Обновляю память...")
    await auto_update_memory(load_history())
    await update.message.reply_text("✅ Готово.")


async def cmd_model(update: Update, ctx: ContextTypes.DEFAULT_TYPE):
    if not is_allowed(update): return
    global MODEL
    text = update.message.text.replace("/model", "").strip()
    if text:
        MODEL = text
        await update.message.reply_text(f"✅ Модель: {MODEL}")
    else:
        await update.message.reply_text(f"Текущая модель: `{MODEL}`")


async def handle_message(update: Update, ctx: ContextTypes.DEFAULT_TYPE):
    if not is_allowed(update): return
    user_text = update.message.text
    if not user_text: return
    await update.message.chat.send_action("typing")
    reply = await ask_model(user_text)
    if len(reply) > 4096:
        for i in range(0, len(reply), 4096):
            await update.message.reply_text(reply[i:i+4096])
    else:
        await update.message.reply_text(reply)


async def handle_photo(update: Update, ctx: ContextTypes.DEFAULT_TYPE):
    if not is_allowed(update): return

    await update.message.chat.send_action("typing")

    # Скачиваем фото в память
    photo = update.message.photo[-1]  # берём максимальное разрешение
    caption = update.message.caption or "Что на этом фото? Если есть текст — переведи на русский."

    try:
        file = await ctx.bot.get_file(photo.file_id)
        import io, base64
        buf = io.BytesIO()
        await file.download_to_memory(buf)
        b64 = base64.b64encode(buf.getvalue()).decode("utf-8")

        history = load_history()
        system = load_system_prompt()

        messages = [{"role": "system", "content": system}] + history + [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": caption},
                    {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{b64}"}},
                ],
            }
        ]

        r = await client.chat.completions.create(
            model=MODEL, messages=messages, max_tokens=1500, temperature=0.7,
        )
        reply = r.choices[0].message.content or "(пустой ответ)"

        # Сохраняем в историю
        history.append({"role": "user", "content": f"[фото] {caption}"})
        history.append({"role": "assistant", "content": reply})
        save_history(history)

    except Exception as e:
        log.error("Ошибка обработки фото: %s", e)
        reply = f"⚠️ Не удалось обработать фото: {e}"

    if len(reply) > 4096:
        for i in range(0, len(reply), 4096):
            await update.message.reply_text(reply[i:i+4096])
    else:
        await update.message.reply_text(reply)


# ─────────────────────────────────────────────
# ЗАПУСК
# ─────────────────────────────────────────────
def main():
    log.info("Запуск Cosmo ✨ (модель: %s, авто-память: каждые %d сообщ.)", MODEL, AUTO_MEMORY_INTERVAL)

    app = Application.builder().token(TELEGRAM_TOKEN).build()
    app.add_handler(CommandHandler("start",    cmd_start))
    app.add_handler(CommandHandler("clear",    cmd_clear))
    app.add_handler(CommandHandler("remember", cmd_remember))
    app.add_handler(CommandHandler("memory",   cmd_memory))
    app.add_handler(CommandHandler("diary",    cmd_diary))
    app.add_handler(CommandHandler("memorize", cmd_memorize))
    app.add_handler(CommandHandler("model",    cmd_model))
    app.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, handle_message))
    app.add_handler(MessageHandler(filters.PHOTO, handle_photo))
    app.run_polling(drop_pending_updates=True)


if __name__ == "__main__":
    main()