import os
import json
import ast
import math
import time
import asyncio
import threading
from collections import defaultdict, deque

import wikipedia
import torch
from telegram import Update, InlineKeyboardButton, InlineKeyboardMarkup
from telegram.constants import ChatAction
from telegram.ext import (
    ApplicationBuilder,
    MessageHandler,
    CommandHandler,
    CallbackQueryHandler,
    ContextTypes,
    filters,
)
from transformers import AutoTokenizer, AutoModelForCausalLM

MODEL_NAME = "Qwen/Qwen2-1.5B-Instruct"
MAX_HISTORY = 12
MAX_STEPS = 4
MAX_NEW_TOKENS_JSON = 220

torch.set_num_threads(max(1, os.cpu_count() // 2))

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.float32,
    device_map="cpu",
    low_cpu_mem_usage=True
)
model.eval()

memory = defaultdict(lambda: deque(maxlen=MAX_HISTORY))
button_state = defaultdict(dict)

PLANNER_SYSTEM_PROMPT = """
You are an advanced Telegram AI agent.

You must never expose internal tool calls, JSON planning, scratch work, or control tokens to the user.

You have exactly three response formats.

1) Final:
{"type":"final","text":"your message to the user"}

2) Buttons:
{"type":"buttons","text":"question for the user","buttons":[{"id":"choice_1","label":"Yes"},{"id":"choice_2","label":"No"}]}

3) Tool:
{"type":"tool","name":"wiki_search","arguments":{"query":"Finland"}}
{"type":"tool","name":"calculate","arguments":{"expression":"(25*17)/5"}}

Rules:
- Output exactly one JSON object.
- No markdown fences.
- No extra text before or after JSON.
- The button type must be exactly "buttons".
- Use buttons when the user should choose between a few short options.
- Use wiki_search for factual topics, places, people, concepts, summaries.
- Use calculate for arithmetic or formula evaluation.
- If you are unsure, do not call a tool. Respond with a final answer instead.
- After receiving a tool result, continue and respond with exactly one JSON object.
- Prefer Finnish if the user speaks Finnish.
- Never output tokens like <|end|>, <|im_start|>, <|im_end|>.
"""

def safe_calculate(expression: str) -> str:
    allowed_names = {
        "abs": abs,
        "round": round,
        "min": min,
        "max": max,
        "pow": pow,
        "sqrt": math.sqrt,
        "sin": math.sin,
        "cos": math.cos,
        "tan": math.tan,
        "pi": math.pi,
        "e": math.e,
    }

    allowed_nodes = (
        ast.Expression,
        ast.BinOp,
        ast.UnaryOp,
        ast.Num,
        ast.Constant,
        ast.Add,
        ast.Sub,
        ast.Mult,
        ast.Div,
        ast.FloorDiv,
        ast.Mod,
        ast.Pow,
        ast.USub,
        ast.UAdd,
        ast.Load,
        ast.Call,
        ast.Name,
        ast.Tuple,
        ast.List,
    )

    try:
        tree = ast.parse(expression, mode="eval")
    except Exception:
        return "Virhe: laskua ei voitu lukea."

    for node in ast.walk(tree):
        if not isinstance(node, allowed_nodes):
            return "Virhe: laskua ei voitu suorittaa turvallisesti."
        if isinstance(node, ast.Call):
            if not isinstance(node.func, ast.Name):
                return "Virhe: laskua ei voitu suorittaa turvallisesti."
            if node.func.id not in allowed_names:
                return "Virhe: laskua ei voitu suorittaa turvallisesti."
        if isinstance(node, ast.Name):
            if node.id not in allowed_names:
                return "Virhe: laskua ei voitu suorittaa turvallisesti."

    try:
        result = eval(compile(tree, "<expr>", "eval"), {"__builtins__": {}}, allowed_names)
        return str(result)
    except Exception:
        return "Virhe: laskua ei voitu suorittaa."

def wiki_search(query: str) -> str:
    query = query.strip()
    if not query:
        return "Virhe: tyhjä hakukysely."

    try:
        wikipedia.set_lang("fi")
        try:
            return wikipedia.summary(query, sentences=3, auto_suggest=True)
        except Exception:
            page = wikipedia.page(query, auto_suggest=True)
            return page.summary[:1200]
    except Exception:
        try:
            wikipedia.set_lang("en")
            try:
                return wikipedia.summary(query, sentences=3, auto_suggest=True)
            except Exception:
                page = wikipedia.page(query, auto_suggest=True)
                return page.summary[:1200]
        except Exception:
            return f"En löytänyt hakutulosta haulle: {query}"

TOOLS = {
    "wiki_search": wiki_search,
    "calculate": safe_calculate,
}

def clean_text(text: str) -> str:
    bad_tokens = [
        "<|end|>",
        "<|im_start|>",
        "<|im_end|>",
        "<|assistant|>",
        "<|user|>",
        "<|system|>",
        "</s>",
    ]
    for token in bad_tokens:
        text = text.replace(token, "")
    return text.strip()

def extract_json(text: str):
    text = text.strip()
    decoder = json.JSONDecoder()
    for i, ch in enumerate(text):
        if ch == "{":
            try:
                obj, end = decoder.raw_decode(text[i:])
                trailing = text[i + end:].strip()
                if trailing:
                    continue
                return obj
            except Exception:
                continue
    return None

def build_planner_messages(user_id: int, user_text: str):
    messages = [{"role": "system", "content": PLANNER_SYSTEM_PROMPT}]
    for item in memory[user_id]:
        messages.append(item)
    messages.append({"role": "user", "content": user_text})
    return messages

def generate_chat_text(messages, max_new_tokens=220):
    prompt = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )
    inputs = tokenizer(prompt, return_tensors="pt")
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=False,
            use_cache=True,
            pad_token_id=tokenizer.eos_token_id
        )
    new_tokens = outputs[0][inputs["input_ids"].shape[1]:]
    text = tokenizer.decode(new_tokens, skip_special_tokens=False)
    return clean_text(text)

def run_agent(user_id: int, user_text: str):
    messages = build_planner_messages(user_id, user_text)

    for _ in range(MAX_STEPS):
        raw = generate_chat_text(messages, max_new_tokens=MAX_NEW_TOKENS_JSON)
        data = extract_json(raw)

        if not isinstance(data, dict):
            messages.append({"role": "assistant", "content": raw})
            messages.append({
                "role": "user",
                "content": 'Your previous response was invalid. Output ONLY one valid JSON object.'
            })
            continue

        response_type = data.get("type")

        if response_type == "final":
            text = clean_text(str(data.get("text", "")))
            if not text:
                text = "En saanut muodostettua vastausta."
            return {"type": "final", "text": text}

        if response_type == "buttons":
            text = clean_text(str(data.get("text", "")))
            buttons = data.get("buttons", [])
            normalized = []
            if isinstance(buttons, list):
                for b in buttons[:6]:
                    if isinstance(b, dict):
                        bid = str(b.get("id", "")).strip()
                        label = str(b.get("label", "")).strip()
                        if bid and label:
                            normalized.append({"id": bid[:32], "label": label[:40]})
            if text and normalized:
                return {"type": "buttons", "text": text, "buttons": normalized}
            messages.append({"role": "assistant", "content": json.dumps(data, ensure_ascii=False)})
            messages.append({
                "role": "user",
                "content": 'That buttons response was invalid. Output ONLY one valid JSON object.'
            })
            continue

        if response_type == "tool":
            name = data.get("name")
            arguments = data.get("arguments", {})

            if name not in TOOLS:
                return {"type": "final", "text": f"Tuntematon työkalu: {name}"}

            if not isinstance(arguments, dict):
                return {"type": "final", "text": "Työkalun argumentit olivat virheelliset."}

            if name == "wiki_search":
                query = str(arguments.get("query", "")).strip()
                result = TOOLS[name](query)
            elif name == "calculate":
                expression = str(arguments.get("expression", "")).strip()
                result = TOOLS[name](expression)
            else:
                result = "Työkalua ei voitu suorittaa."

            messages.append({"role": "assistant", "content": json.dumps(data, ensure_ascii=False)})
            messages.append({
                "role": "user",
                "content": f"Tool result for {name}:\n{result}\nNow continue and respond with exactly one JSON object."
            })
            continue

        messages.append({"role": "assistant", "content": json.dumps(data, ensure_ascii=False)})
        messages.append({
            "role": "user",
            "content": 'That response type was invalid. Output ONLY one valid JSON object.'
        })

    return {"type": "final", "text": "Pyyntö vaati liikaa välivaiheita."}

async def typing_loop(chat, stop_event: asyncio.Event):
    while not stop_event.is_set():
        try:
            await chat.send_action(ChatAction.TYPING)
        except Exception:
            return
        try:
            await asyncio.wait_for(stop_event.wait(), timeout=2.0)
        except asyncio.TimeoutError:
            pass

def chunk_text_for_stream(text: str):
    words = text.split()
    if not words:
        return [""]

    chunks = []
    current = ""

    for word in words:
        candidate = f"{current} {word}".strip()
        if len(candidate) >= 35:
            chunks.append(candidate)
            current = ""
        else:
            current = candidate

    if current:
        chunks.append(current)

    return chunks

async def stream_text_reply(message, text: str):
    text = clean_text(text)
    if not text:
        text = " "

    chunks = chunk_text_for_stream(text)
    sent = await message.reply_text("...")
    assembled = ""
    last_edit_time = 0.0

    for i, chunk in enumerate(chunks):
        assembled = f"{assembled} {chunk}".strip()
        now = time.time()

        if i < len(chunks) - 1:
            if now - last_edit_time < 0.55:
                await asyncio.sleep(0.55 - (now - last_edit_time))

        safe_text = assembled[:4096]
        try:
            await sent.edit_text(safe_text)
            last_edit_time = time.time()
        except Exception:
            pass

    return sent

async def process_user_text(message, context: ContextTypes.DEFAULT_TYPE, user_id: int, text: str):
    stop_event = asyncio.Event()
    typing_task = asyncio.create_task(typing_loop(message.chat, stop_event))

    try:
        result = await asyncio.to_thread(run_agent, user_id, text)
    finally:
        stop_event.set()
        await typing_task

    memory[user_id].append({"role": "user", "content": text})

    if result["type"] == "buttons":
        keyboard = []
        button_state[user_id] = {}

        for b in result["buttons"]:
            button_state[user_id][b["id"]] = b["label"]
            keyboard.append([InlineKeyboardButton(b["label"], callback_data=f"btn:{b['id']}")])

        memory[user_id].append({
            "role": "assistant",
            "content": json.dumps(
                {"type": "buttons", "text": result["text"], "buttons": result["buttons"]},
                ensure_ascii=False
            )
        })

        await message.reply_text(
            result["text"],
            reply_markup=InlineKeyboardMarkup(keyboard)
        )
        return

    reply = clean_text(result["text"])
    memory[user_id].append({
        "role": "assistant",
        "content": json.dumps({"type": "final", "text": reply}, ensure_ascii=False)
    })

    await stream_text_reply(message, reply)

async def handle_message(update: Update, context: ContextTypes.DEFAULT_TYPE):
    if not update.message or not update.effective_user:
        return

    text = (update.message.text or "").strip()
    if not text:
        return

    await process_user_text(update.message, context, update.effective_user.id, text)

async def handle_button(update: Update, context: ContextTypes.DEFAULT_TYPE):
    query = update.callback_query
    if not query or not update.effective_user:
        return

    await query.answer()

    data = query.data or ""
    if not data.startswith("btn:"):
        return

    button_id = data[4:]
    user_id = update.effective_user.id
    label = button_state[user_id].get(button_id)

    if not label:
        await query.message.reply_text("Tämä valinta ei ole enää voimassa.")
        return

    await process_user_text(query.message, context, user_id, label)

async def start(update: Update, context: ContextTypes.DEFAULT_TYPE):
    if not update.message:
        return
    await update.message.reply_text("Moi. Olen AI-agentti. Laita viestiä.")

async def reset_chat(update: Update, context: ContextTypes.DEFAULT_TYPE):
    if not update.message or not update.effective_user:
        return
    user_id = update.effective_user.id
    memory[user_id].clear()
    button_state[user_id].clear()
    await update.message.reply_text("Muisti nollattu.")

def main():
    token = os.environ["TELEGRAM_TOKEN"]

    app = ApplicationBuilder().token(token).build()

    app.add_handler(CommandHandler("start", start))
    app.add_handler(CommandHandler("reset", reset_chat))
    app.add_handler(CallbackQueryHandler(handle_button))
    app.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, handle_message))

    app.run_polling(drop_pending_updates=True)

if __name__ == "__main__":
    main()