"""
FORGE Telegram Bot
==================
Full async Telegram bot powered by python-telegram-bot v21.
- /generate <prompt> — start app generation
- /status — check active session
- /cancel — cancel active session
- /config — show current settings
- Sends live step-by-step progress messages
- Sends zip file on completion
- Handles clarifying questions interactively
- Rate limiting per chat
"""
from __future__ import annotations

import asyncio
import logging
import time
import io
from pathlib import Path
from typing import Optional
from dataclasses import dataclass, field

from forge.core.session import GenerationSession, LogEntry, AgentRun, AgentStatus
from forge.core.pipeline import Pipeline
from forge.settings import Settings

logger = logging.getLogger("forge.telegram")

# Import guard — python-telegram-bot is optional at import time
try:
    from telegram import (
        Update,
        Document,
        InlineKeyboardButton,
        InlineKeyboardMarkup,
        ParseMode,
    )
    from telegram.ext import (
        Application,
        CommandHandler,
        MessageHandler,
        CallbackQueryHandler,
        ContextTypes,
        filters,
        ConversationHandler,
    )
    from telegram.constants import ParseMode, ChatAction
    from telegram.error import TelegramError
    TELEGRAM_AVAILABLE = True
except ImportError:
    TELEGRAM_AVAILABLE = False


# ── Conversation states ───────────────────────────────────────────────────────
AWAITING_PROMPT = 0
AWAITING_CLARIFICATION = 1
AWAITING_COMPLEXITY = 2


@dataclass
class ChatSession:
    """Per-chat active state."""
    chat_id: int
    gen_session: Optional[GenerationSession] = None
    status_message_id: Optional[int] = None
    started_at: float = field(default_factory=time.time)
    is_busy: bool = False
    pending_questions: list[str] = field(default_factory=list)
    answers_so_far: dict[str, str] = field(default_factory=dict)
    current_question_idx: int = 0


class ForgeBot:
    """
    Full-featured async Telegram bot for FORGE.
    Each chat gets its own isolated session.
    """

    # Max chars per Telegram message
    MAX_MSG_LEN = 4096
    # Progress update throttle (seconds)
    PROGRESS_THROTTLE = 3.0

    def __init__(self, settings: Settings):
        if not TELEGRAM_AVAILABLE:
            raise RuntimeError(
                "python-telegram-bot not installed. Run: pip install 'forge-ai[telegram]'"
            )
        self.settings = settings
        self.pipeline = Pipeline(settings)
        self._chats: dict[int, ChatSession] = {}
        self._app: Optional["Application"] = None

    # ── Setup ─────────────────────────────────────────────────────────────────

    def build_app(self) -> "Application":
        token = self.settings.telegram_bot_token
        app = Application.builder().token(token).build()

        app.add_handler(CommandHandler("start", self.cmd_start))
        app.add_handler(CommandHandler("help", self.cmd_help))
        app.add_handler(CommandHandler("generate", self.cmd_generate))
        app.add_handler(CommandHandler("status", self.cmd_status))
        app.add_handler(CommandHandler("cancel", self.cmd_cancel))
        app.add_handler(CommandHandler("config", self.cmd_config))
        app.add_handler(CommandHandler("doctor", self.cmd_doctor))
        app.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, self.handle_text))
        app.add_error_handler(self.error_handler)

        self._app = app
        return app

    def run(self) -> None:
        """Start polling (blocking)."""
        app = self.build_app()
        logger.info("FORGE Telegram Bot starting…")
        app.run_polling(allowed_updates=Update.ALL_TYPES)

    # ── Chat session helpers ──────────────────────────────────────────────────

    def _get_chat(self, chat_id: int) -> ChatSession:
        if chat_id not in self._chats:
            self._chats[chat_id] = ChatSession(chat_id=chat_id)
        return self._chats[chat_id]

    async def _safe_send(self, context, chat_id: int, text: str, **kwargs) -> Optional[int]:
        """Send a message, splitting if over Telegram's limit."""
        chunks = [text[i:i+self.MAX_MSG_LEN] for i in range(0, len(text), self.MAX_MSG_LEN)]
        msg_id = None
        for chunk in chunks:
            try:
                msg = await context.bot.send_message(chat_id=chat_id, text=chunk, **kwargs)
                msg_id = msg.message_id
            except TelegramError as e:
                logger.warning(f"send_message failed: {e}")
        return msg_id

    async def _safe_edit(self, context, chat_id: int, message_id: int, text: str) -> None:
        try:
            await context.bot.edit_message_text(
                chat_id=chat_id,
                message_id=message_id,
                text=text[:self.MAX_MSG_LEN],
            )
        except TelegramError:
            pass

    # ── Commands ──────────────────────────────────────────────────────────────

    async def cmd_start(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
        text = (
            "🔨 *FORGE AI — Mobile App Generator*\n\n"
            "Turn any app idea into production-ready React Native + Expo code.\n\n"
            "*Commands:*\n"
            "`/generate <your app idea>` — Start generating\n"
            "`/status` — Check active generation\n"
            "`/cancel` — Cancel current generation\n"
            "`/config` — View current settings\n"
            "`/doctor` — Test API connections\n"
            "`/help` — Full help\n\n"
            "_Powered by Groq LLaMA 3.3 + NVIDIA NIM CodeLlama_"
        )
        await update.message.reply_text(text, parse_mode=ParseMode.MARKDOWN)

    async def cmd_help(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
        text = (
            "🔨 *FORGE Help*\n\n"
            "*Generate an app:*\n"
            "`/generate A habit tracker with streaks and reminders`\n\n"
            "*What happens:*\n"
            "1️⃣ PLANNER analyzes your prompt\n"
            "2️⃣ ARCHITECT designs file structure\n"
            "3️⃣ CODE GEN writes every file (NVIDIA NIM)\n"
            "4️⃣ VALIDATOR fixes critical issues\n"
            "5️⃣ TESTER writes Jest test suites\n"
            "6️⃣ PACKAGER zips and sends to you\n\n"
            "*Output:*\n"
            "• Complete React Native + Expo project\n"
            "• TypeScript throughout\n"
            "• Jest test suite\n"
            "• GitHub Actions CI/CD\n"
            "• EAS build config\n"
            "• `npm install && npx expo start` — done\n\n"
            "*Complexity settings (set via* `forge config`*):\n*"
            "• `minimal` — 4 screens\n"
            "• `standard` — 7 screens _(default)_\n"
            "• `full` — 12 screens\n"
            "• `enterprise` — 18 screens"
        )
        await update.message.reply_text(text, parse_mode=ParseMode.MARKDOWN)

    async def cmd_generate(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
        chat_id = update.effective_chat.id
        chat = self._get_chat(chat_id)

        if chat.is_busy:
            await update.message.reply_text(
                "⏳ Already generating an app. Use /status to check, or /cancel to stop."
            )
            return

        prompt = " ".join(context.args).strip() if context.args else ""
        if not prompt:
            await update.message.reply_text(
                "Please include your app idea:\n"
                "`/generate A recipe app with meal planning and shopping lists`",
                parse_mode=ParseMode.MARKDOWN,
            )
            return

        if len(prompt) < 10:
            await update.message.reply_text(
                "⚠️ Please describe your app in more detail (at least 10 characters)."
            )
            return

        await self._start_generation(update, context, chat, prompt)

    async def _start_generation(
        self,
        update: Update,
        context: ContextTypes.DEFAULT_TYPE,
        chat: ChatSession,
        prompt: str,
    ) -> None:
        chat.is_busy = True
        chat.started_at = time.time()

        # Send initial status message
        status_msg = await update.message.reply_text(
            f"🚀 *Starting FORGE pipeline…*\n\n"
            f"📝 Prompt: `{prompt[:100]}`\n"
            f"⚙️ Complexity: `{self.settings.complexity}`\n\n"
            f"_Initializing agents…_",
            parse_mode=ParseMode.MARKDOWN,
        )
        chat.status_message_id = status_msg.message_id

        # Log buffer for live updates
        log_buffer: list[str] = []
        last_update_time = [0.0]
        agent_states: dict[str, str] = {}

        def format_status() -> str:
            agent_icons = {
                "PLANNER": "🧠", "ARCHITECT": "🏗", "CODEGEN": "⚙️",
                "VALIDATOR": "✅", "TESTER": "🧪", "PACKAGER": "📦", "FORGE": "🔨",
            }
            lines = [f"🔨 *FORGE — Generating* `{prompt[:60]}`\n"]

            for agent, state in agent_states.items():
                icon = agent_icons.get(agent, "•")
                lines.append(f"{icon} `{agent}`: {state}")

            if log_buffer:
                lines.append("\n*Latest:*")
                lines.extend(f"`{l}`" for l in log_buffer[-5:])

            elapsed = time.time() - chat.started_at
            lines.append(f"\n⏱ `{elapsed:.0f}s elapsed`")
            return "\n".join(lines)

        async def on_log(entry: LogEntry) -> None:
            level_prefix = {"success": "✓", "warning": "⚠", "error": "✗", "highlight": "★"}.get(
                entry.level, "›"
            )
            log_buffer.append(f"{level_prefix} [{entry.agent}] {entry.message[:80]}")
            if len(log_buffer) > 30:
                log_buffer.pop(0)

            # Throttle edits
            now = time.time()
            if now - last_update_time[0] > self.PROGRESS_THROTTLE:
                last_update_time[0] = now
                await self._safe_edit(
                    context,
                    chat.chat_id,
                    chat.status_message_id,
                    format_status(),
                )

        def on_log_sync(entry: LogEntry) -> None:
            asyncio.create_task(on_log(entry))

        def on_agent_update(run: AgentRun) -> None:
            status_map = {
                AgentStatus.RUNNING: "⏳ running…",
                AgentStatus.DONE: f"✅ done ({run.elapsed:.1f}s)",
                AgentStatus.FAILED: f"❌ failed",
                AgentStatus.PENDING: "⬜ waiting",
            }
            agent_states[run.name] = status_map.get(run.status, "")

        def on_progress(current: int, total: int, path: str) -> None:
            agent_states["CODEGEN"] = f"⏳ {current}/{total} — {Path(path).name}"

        try:
            session = await self.pipeline.run(
                prompt=prompt,
                clarification_answers=chat.answers_so_far or None,
                on_log=on_log_sync,
                on_agent_update=on_agent_update,
                on_progress=on_progress,
            )

            chat.gen_session = session

            # Handle clarification needed
            if session.needs_clarification and not chat.answers_so_far:
                chat.pending_questions = session.plan.clarifying_questions
                chat.is_busy = False
                await self._ask_clarification(update, context, chat)
                return

            # Success — send results
            await self._send_results(context, chat, session)

        except Exception as e:
            logger.exception("Generation failed")
            await self._safe_send(
                context,
                chat.chat_id,
                f"❌ *Generation failed*\n\n`{type(e).__name__}: {str(e)[:300]}`\n\nPlease try again.",
                parse_mode=ParseMode.MARKDOWN,
            )
        finally:
            chat.is_busy = False
            chat.answers_so_far = {}
            chat.pending_questions = []
            chat.current_question_idx = 0

    async def _ask_clarification(
        self,
        update: Update,
        context: ContextTypes.DEFAULT_TYPE,
        chat: ChatSession,
    ) -> None:
        if not chat.pending_questions:
            return

        q = chat.pending_questions[chat.current_question_idx]
        total = len(chat.pending_questions)
        idx = chat.current_question_idx

        await context.bot.send_message(
            chat_id=chat.chat_id,
            text=(
                f"🤔 *Clarification needed ({idx+1}/{total}):*\n\n"
                f"{q}\n\n"
                f"_Reply with your answer, or send_ `/skip` _to use defaults._"
            ),
            parse_mode=ParseMode.MARKDOWN,
        )

    async def handle_text(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
        chat_id = update.effective_chat.id
        chat = self._get_chat(chat_id)
        text = update.message.text.strip()

        # If awaiting clarification answers
        if chat.pending_questions and not chat.is_busy:
            q = chat.pending_questions[chat.current_question_idx]
            chat.answers_so_far[q] = text
            chat.current_question_idx += 1

            if chat.current_question_idx >= len(chat.pending_questions):
                # All answered — restart generation with answers
                await update.message.reply_text("✅ Got it! Resuming generation…")
                prompt = chat.gen_session.prompt if chat.gen_session else text
                await self._start_generation(update, context, chat, prompt)
            else:
                await self._ask_clarification(update, context, chat)
            return

        # Otherwise treat as new prompt
        if not chat.is_busy and len(text) > 10:
            await update.message.reply_text(
                f"💡 To generate an app, use:\n`/generate {text}`",
                parse_mode=ParseMode.MARKDOWN,
            )

    async def _send_results(
        self,
        context: ContextTypes.DEFAULT_TYPE,
        chat: ChatSession,
        session: GenerationSession,
    ) -> None:
        plan = session.plan
        elapsed = session.elapsed

        # Summary message
        summary = (
            f"✅ *{plan.app_name}* — Generation Complete!\n\n"
            f"_{plan.tagline}_\n\n"
            f"📱 *{len(plan.screens)} screens* | "
            f"📄 *{session.total_files} files* | "
            f"📝 *{session.total_lines:,} lines*\n"
            f"⏱ Generated in `{elapsed:.1f}s`\n\n"
            f"*Screens:*\n"
        )
        for s in plan.screens[:8]:
            summary += f"• `{s['name']}` — {s['purpose']}\n"

        summary += (
            f"\n*Tech Stack:*\n"
            f"• React Native + Expo SDK 50\n"
            f"• TypeScript (strict)\n"
            f"• Navigation: {plan.navigation}\n"
            f"• Storage: {plan.storage}\n"
            f"• Tests: Jest + RNTL\n"
            f"• CI/CD: GitHub Actions + EAS\n\n"
            f"*Next steps:*\n"
            f"```\n"
            f"cd {plan.slug}\n"
            f"npm install\n"
            f"npx expo start\n"
            f"```"
        )

        await self._safe_send(
            context,
            chat.chat_id,
            summary,
            parse_mode=ParseMode.MARKDOWN,
        )

        # Send zip
        if session.zip_path and session.zip_path.exists():
            zip_size_mb = session.zip_path.stat().st_size / (1024 * 1024)
            await context.bot.send_chat_action(
                chat_id=chat.chat_id, action=ChatAction.UPLOAD_DOCUMENT
            )
            with open(session.zip_path, "rb") as f:
                await context.bot.send_document(
                    chat_id=chat.chat_id,
                    document=f,
                    filename=f"{plan.slug}.zip",
                    caption=(
                        f"📦 *{plan.app_name}* — {session.total_files} files, "
                        f"{zip_size_mb:.1f} MB\n"
                        f"_Unzip → `npm install` → `npx expo start`_"
                    ),
                    parse_mode=ParseMode.MARKDOWN,
                )
        else:
            await self._safe_send(
                context,
                chat.chat_id,
                f"⚠️ Zip not found at {session.zip_path}. Project was written to disk at: `{session.project_dir}`",
            )

    async def cmd_status(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
        chat = self._get_chat(update.effective_chat.id)
        if chat.is_busy:
            elapsed = time.time() - chat.started_at
            await update.message.reply_text(
                f"⏳ Generation in progress ({elapsed:.0f}s elapsed…)"
            )
        elif chat.gen_session:
            s = chat.gen_session
            await update.message.reply_text(
                f"✅ Last generation: *{s.plan.app_name if s.plan else 'unknown'}*\n"
                f"Files: {s.total_files} | Lines: {s.total_lines:,} | "
                f"Time: {s.elapsed:.1f}s",
                parse_mode=ParseMode.MARKDOWN,
            )
        else:
            await update.message.reply_text(
                "No active session. Use `/generate <prompt>` to start.",
                parse_mode=ParseMode.MARKDOWN,
            )

    async def cmd_cancel(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
        chat = self._get_chat(update.effective_chat.id)
        if chat.is_busy:
            chat.is_busy = False
            await update.message.reply_text("🛑 Cancellation requested. (Note: ongoing API calls may complete.)")
        else:
            await update.message.reply_text("Nothing to cancel.")

    async def cmd_config(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
        s = self.settings
        text = (
            "⚙️ *FORGE Configuration*\n\n"
            f"Plan model: `{s.plan_model}`\n"
            f"Code model: `{s.code_model}`\n"
            f"Complexity: `{s.complexity}`\n"
            f"Framework: `{s.framework}`\n"
            f"TypeScript: `{s.typescript}`\n"
            f"Include tests: `{s.include_tests}`\n"
            f"Include CI: `{s.include_ci}`\n"
            f"Output: `{s.output_dir}`\n\n"
            f"Groq key: `{s.mask(s.groq_api_key)}`\n"
            f"NVIDIA key: `{s.mask(s.nvidia_api_key)}`"
        )
        await update.message.reply_text(text, parse_mode=ParseMode.MARKDOWN)

    async def cmd_doctor(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
        await update.message.reply_text("🔍 Testing API connections…")
        results = await self.pipeline.validate_keys()
        groq_ok = "✅" if results["groq"] else "❌"
        nvidia_ok = "✅" if results["nvidia"] else "❌"
        text = (
            f"*API Health Check*\n\n"
            f"{groq_ok} Groq API: `{'OK' if results['groq'] else 'FAILED'}`\n"
            f"{nvidia_ok} NVIDIA NIM: `{'OK' if results['nvidia'] else 'FAILED'}`"
        )
        await update.message.reply_text(text, parse_mode=ParseMode.MARKDOWN)

    async def error_handler(self, update: object, context: ContextTypes.DEFAULT_TYPE) -> None:
        logger.error(f"Telegram error: {context.error}", exc_info=context.error)
