import asyncio
import json
import logging
from typing import Any

from aiogram import Bot, Dispatcher, types
from aiogram.enums import ParseMode

logger = logging.getLogger("cvc.telegram_handler")

# We will maintain a global dict of pending tool approvals for Telegram
# message_id -> { "tool": tc, "event": asyncio.Event(), "decision": "" }
pending_approvals = {}

# Locks to prevent concurrent execution for the same user
session_locks = {}

def register_handlers(dp: Dispatcher, bot: Bot, settings: Any, workspace_mgr: Any):
    
    @dp.message()
    async def handle_message(message: types.Message):
        if message.from_user.id not in settings.telegram.allowlist:
            logger.warning(f"Unauthorized access from Telegram ID: {message.from_user.id}")
            return
            
        import cvc.gateway as gw
        
        if not gw._agent_llm or not gw._agent_tools:
            await message.reply("CVC Agent LLM is not initialized yet.")
            return
            
        session_key = f"{message.chat.id}_{message.message_thread_id or 0}"
        
        # Ensure only one task runs per session at a time
        if session_key not in session_locks:
            session_locks[session_key] = asyncio.Lock()
            
        if session_locks[session_key].locked():
            await message.reply("⚠️ <i>A task is currently running. Please wait until it completes.</i>", parse_mode=ParseMode.HTML)
            return
            
        async with session_locks[session_key]:
            # Push message to LLM context
            if not hasattr(gw, "_tg_sessions"):
                gw._tg_sessions = {}
                
            if session_key not in gw._tg_sessions:
                gw._tg_sessions[session_key] = [{"role": "system", "content": gw._system_prompt}]
                
            llm_messages = gw._tg_sessions[session_key]
            llm_messages.append({"role": "user", "content": message.text or ""})
            
            # Start streaming response
            from cvc.adapters.telegram_streamer import TelegramStreamer
            
            reply_msg = await message.reply("⚙️ <i>Thinking...</i>", parse_mode=ParseMode.HTML)
            streamer = TelegramStreamer(bot, message.chat.id, reply_msg.message_id)
            
            # Run agentic loop
            await run_agentic_loop(bot, message, streamer, llm_messages, gw, session_key)

    @dp.callback_query()
    async def handle_callback(callback: types.CallbackQuery):
        if callback.from_user.id not in settings.telegram.allowlist:
            return
            
        data = callback.data
        msg_id = callback.message.message_id
        
        if msg_id in pending_approvals:
            approval = pending_approvals[msg_id]
            approval["decision"] = data
            approval["event"].set()
            
            # Edit the message to show decision
            await callback.message.edit_text(
                f"{callback.message.text}\n\n<b>Decision:</b> {data}",
                parse_mode=ParseMode.HTML
            )
            await callback.answer(f"Decision registered: {data}")
        else:
            await callback.answer("This approval request has expired or is invalid.", show_alert=True)

async def run_agentic_loop(bot: Bot, message: types.Message, streamer, llm_messages: list, gw, session_key: str):
    iteration = 0
    max_iterations = getattr(gw, "_MAX_AGENT_ITERATIONS", 50)
    
    try:
        while iteration < max_iterations:
            iteration += 1
            response_text = ""
            tool_calls = []
            tool_call_args_buffer = {}
            gemini_parts = None
            
            temp = 0.7 if iteration == 1 else 0.5
            max_tok = 8192 if iteration == 1 else 16384
            
            async for event in gw._agent_llm.chat_stream(
                messages=llm_messages,
                tools=gw._agent_tools,
                temperature=temp,
                max_tokens=max_tok,
            ):
                if event.type == "text_delta" and event.text:
                    response_text += event.text
                    streamer.push(event.text)
                    
                elif event.type == "tool_call_start" and event.tool_call:
                    tc = event.tool_call
                    tool_calls.append(tc)
                    tool_call_args_buffer[event.tool_call_index] = ""
                    
                elif event.type == "tool_call_delta":
                    idx = event.tool_call_index
                    if idx in tool_call_args_buffer:
                        tool_call_args_buffer[idx] += event.args_delta
                        if tool_calls and idx < len(tool_calls):
                            try:
                                tool_calls[idx].arguments = json.loads(tool_call_args_buffer[idx])
                            except Exception:
                                pass
                                
                elif event.type == "done":
                    if event._provider_meta.get("gemini_parts"):
                        gemini_parts = event._provider_meta["gemini_parts"]

            # Output streaming stops when done with text
            await streamer.finalize()
            
            if not tool_calls:
                if response_text:
                    llm_messages.append({"role": "assistant", "content": response_text})
                break
                
            assistant_msg = {"role": "assistant", "content": response_text or ""}
            if gw._agent_llm.provider in ("openai", "ollama", "lmstudio", "google"):
                assistant_msg["tool_calls"] = [
                    {"id": tc.id, "type": "function", "function": {"name": tc.name, "arguments": json.dumps(tc.arguments)}}
                    for tc in tool_calls
                ]
                if gemini_parts:
                    assistant_msg["_gemini_parts"] = gemini_parts
            elif gw._agent_llm.provider == "anthropic":
                content_blocks = []
                if response_text:
                    content_blocks.append({"type": "text", "text": response_text})
                for tc in tool_calls:
                    content_blocks.append({"type": "tool_use", "id": tc.id, "name": tc.name, "input": tc.arguments})
                assistant_msg = {"role": "assistant", "content": content_blocks}
                
            llm_messages.append(assistant_msg)
            
            for tc in tool_calls:
                # Permission logic
                needs_confirm = (
                    not gw._autopilot
                    and not getattr(gw, "_session_trust_all", False)
                    and tc.name not in getattr(gw, "_SAFE_TOOLS", set())
                    and tc.name not in getattr(gw, "_session_trusted_tools", set())
                )
                
                decision = "allow_once"
                if needs_confirm:
                    # Ask via Inline Keyboard
                    keyboard = types.InlineKeyboardMarkup(inline_keyboard=[
                        [
                            types.InlineKeyboardButton(text="✅ Approve", callback_data="allow_once"),
                            types.InlineKeyboardButton(text="✅ Always", callback_data="allow_always"),
                        ],
                        [
                            types.InlineKeyboardButton(text="❌ Deny", callback_data="deny"),
                        ]
                    ])
                    args_disp = json.dumps(tc.arguments, indent=2)[:500]
                    prompt_text = f"⚠️ <b>Tool Execution Requested</b>\n\n<b>Tool:</b> <code>{tc.name}</code>\n<b>Args:</b>\n<code>{args_disp}</code>"
                    
                    req_msg = await bot.send_message(
                        chat_id=message.chat.id,
                        message_thread_id=message.message_thread_id,
                        text=prompt_text,
                        reply_markup=keyboard,
                        parse_mode=ParseMode.HTML
                    )
                    
                    evt = asyncio.Event()
                    pending_approvals[req_msg.message_id] = {"event": evt, "decision": "deny"}
                    
                    try:
                        await asyncio.wait_for(evt.wait(), timeout=120.0)
                    except asyncio.TimeoutError:
                        pass
                        
                    decision = pending_approvals[req_msg.message_id]["decision"]
                    del pending_approvals[req_msg.message_id]
                    
                if decision == "allow_always":
                    gw._session_trusted_tools.add(tc.name)
                elif decision == "deny":
                    result = f"Tool '{tc.name}' was denied by the user via Telegram."
                    llm_messages.append({
                        "role": "tool",
                        "tool_call_id": tc.id,
                        "name": tc.name,
                        "content": result,
                    })
                    continue

                # Execute tool
                try:
                    if tc.name == "cvc_switch_workspace":
                        result = await gw._execute_workspace_switch(tc.arguments)
                    else:
                        loop = asyncio.get_event_loop()
                        result = await loop.run_in_executor(
                            None, gw._tool_executor.execute, tc.name, tc.arguments
                        )
                except Exception as exc:
                    result = f"Error executing {tc.name}: {exc}"
                    
                max_out = getattr(gw, "_MAX_TOOL_OUTPUT_LLM", 15000)
                llm_result = result[:max_out] if len(result) > max_out else result
                
                llm_messages.append({
                    "role": "tool",
                    "tool_call_id": tc.id,
                    "name": tc.name,
                    "content": llm_result,
                })

            # Re-initialize streamer for the next turn
            reply_msg = await message.reply("⚙️ <i>Processing...</i>", parse_mode=ParseMode.HTML)
            streamer = TelegramStreamer(bot, message.chat.id, reply_msg.message_id)

    except Exception as exc:
        logger.error(f"Agentic loop error: {exc}", exc_info=True)
        await bot.send_message(message.chat.id, f"<b>Error:</b> {exc}", parse_mode=ParseMode.HTML)
        
    finally:
        await streamer.finalize()
