"""
Standalone Telegram Bot Worker for Kora AI Ecosystem.
Connects Telegram users directly to Kora Chat, Kora Coder 0.5B, and Kora Image Studio.
Uses httpx with explicit timeouts and token redaction for security and stability.
"""

import os
import sys
import time
import json
import socket
import urllib.parse
import urllib.request
from threading import Thread
from typing import Optional, Dict, Any

# Force IPv4 resolution for outbound requests (prevents IPv6 handshake timeouts on cloud containers)
_old_getaddrinfo = socket.getaddrinfo
def _getaddrinfo_ipv4_only(host, port, family=0, type=0, proto=0, flags=0):
    try:
        return _old_getaddrinfo(host, port, socket.AF_INET, type, proto, flags)
    except Exception:
        return _old_getaddrinfo(host, port, family, type, proto, flags)

socket.getaddrinfo = _getaddrinfo_ipv4_only

if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(encoding="utf-8")

# Kora Gateway Endpoint
KORA_ENDPOINT = os.environ.get("KORA_ENDPOINT", os.environ.get("KORA_CUSTOM_ENDPOINT", "http://127.0.0.1:7860")).rstrip("/")
BOT_TOKEN = os.environ.get("TELEGRAM_BOT_TOKEN", "").strip()
TELEGRAM_API_BASE = os.environ.get("TELEGRAM_API_BASE", os.environ.get("TELEGRAM_PROXY", "https://api.telegram.org")).rstrip("/")

def redact_token(text: str) -> str:
    if not BOT_TOKEN:
        return text
    return text.replace(BOT_TOKEN, "[REDACTED_BOT_TOKEN]")

def send_telegram_request(method: str, payload: Optional[Dict[str, Any]] = None, timeout_sec: float = 10.0) -> Optional[Dict[str, Any]]:
    """
    Reusable Telegram API request helper using httpx (or urllib fallback).
    Supports Kontyra API proxy relay if direct Telegram connectivity is restricted by cloud container firewall.
    Redacts bot token from all logs and handles connection/timeout exceptions gracefully.
    """
    if not BOT_TOKEN:
        return None
        
    headers = {
        "Content-Type": "application/json",
        "User-Agent": "KoraBot/1.0 (Kontyra AI Ecosystem)"
    }
    
    # 1. Try Direct Telegram API
    direct_url = f"{TELEGRAM_API_BASE}/bot{BOT_TOKEN}/{method}"
    try:
        req_data = json.dumps(payload).encode("utf-8") if payload else None
        req = urllib.request.Request(direct_url, data=req_data, headers=headers)
        with urllib.request.urlopen(req, timeout=4.0) as resp:
            return json.loads(resp.read().decode("utf-8"))
    except Exception as direct_err:
        pass

    # 2. Try Kontyra API Outbound Relay Proxy
    relay_url = "https://api.kontyra.name.ng/v1/telegram/request"
    try:
        relay_payload = json.dumps({
            "token": BOT_TOKEN,
            "method": method,
            "payload": payload or {}
        }).encode("utf-8")
        req = urllib.request.Request(relay_url, data=relay_payload, headers=headers)
        with urllib.request.urlopen(req, timeout=10.0) as resp:
            res_json = json.loads(resp.read().decode("utf-8"))
            if res_json.get("ok"):
                return res_json
            else:
                return res_json
    except Exception as relay_err:
        safe_err = redact_token(str(relay_err))
        print(f"⚠️ Telegram API request [{method}] relay failed: {safe_err}", flush=True)
        return None

def test_telegram_connectivity() -> bool:
    """
    Startup connectivity diagnostic that tests api.telegram.org without exposing the bot token.
    """
    if not BOT_TOKEN:
        print("ℹ️ Telegram Diagnostic: TELEGRAM_BOT_TOKEN not provided.", flush=True)
        return False

    print("🔍 Testing outbound connectivity to api.telegram.org...", flush=True)
    res = send_telegram_request("getMe", timeout_sec=6.0)
    if res and res.get("ok"):
        bot_username = res.get("result", {}).get("username", "UnknownBot")
        print(f"✅ Telegram Connectivity Check: SUCCESS (Connected as @{bot_username})", flush=True)
        return True
    else:
        print("⚠️ Telegram Connectivity Check: FAILED or TIMED OUT. Gateway will continue operating normally.", flush=True)
        return False

def send_telegram_message(chat_id: int, text: str):
    send_telegram_request("sendMessage", {"chat_id": chat_id, "text": text})

def send_telegram_photo(chat_id: int, photo_path_or_url: str, caption: str = ""):
    if not BOT_TOKEN or not photo_path_or_url:
        return

    # Determine local file and public URL
    local_file = photo_path_or_url
    if not os.path.exists(local_file) and "generated_assets" in photo_path_or_url:
        filename = photo_path_or_url.split("/")[-1]
        potential = os.path.join("public", "generated_assets", filename)
        if os.path.exists(potential):
            local_file = potential

    filename = os.path.basename(local_file) if os.path.exists(local_file) else photo_path_or_url.split("/")[-1]
    public_url = photo_path_or_url if (photo_path_or_url.startswith("http://") or photo_path_or_url.startswith("https://")) else f"https://professorceo-kora-ai.hf.space/public/generated_assets/{filename}"

    # 1. Try sending photo via Public URL (routes via Kontyra API relay proxy if direct Telegram API is blocked)
    print(f"📤 Sending photo to Telegram via public URL: {public_url}", flush=True)
    res = send_telegram_request("sendPhoto", {"chat_id": chat_id, "photo": public_url, "caption": caption}, timeout_sec=15.0)
    if res and res.get("ok"):
        print("✅ Telegram photo delivered successfully via public URL!", flush=True)
        return

    # 2. Try direct binary multipart upload fallback
    if os.path.exists(local_file):
        try:
            import mimetypes, uuid
            boundary = f"----FormBoundary{uuid.uuid4().hex}"
            
            with open(local_file, "rb") as f:
                img_bytes = f.read()
                
            mime_type = mimetypes.guess_type(filename)[0] or "image/png"
            
            body = bytearray()
            body.extend(f"--{boundary}\r\nContent-Disposition: form-data; name=\"chat_id\"\r\n\r\n{chat_id}\r\n".encode("utf-8"))
            if caption:
                body.extend(f"--{boundary}\r\nContent-Disposition: form-data; name=\"caption\"\r\n\r\n{caption}\r\n".encode("utf-8"))
            body.extend(f"--{boundary}\r\nContent-Disposition: form-data; name=\"photo\"; filename=\"{filename}\"\r\nContent-Type: {mime_type}\r\n\r\n".encode("utf-8"))
            body.extend(img_bytes)
            body.extend(f"\r\n--{boundary}--\r\n".encode("utf-8"))
            
            headers = {
                "Content-Type": f"multipart/form-data; boundary={boundary}",
                "User-Agent": "KoraBot/1.0 (Kontyra AI Ecosystem)"
            }
            
            direct_url = f"https://api.telegram.org/bot{BOT_TOKEN}/sendPhoto"
            req = urllib.request.Request(direct_url, data=body, headers=headers)
            with urllib.request.urlopen(req, timeout=25) as resp:
                res_data = json.loads(resp.read().decode("utf-8"))
                if res_data.get("ok"):
                    print("✅ Telegram photo delivered via binary upload!", flush=True)
                    return
        except Exception as e:
            print(f"⚠️ Binary photo upload notice: {e}", flush=True)

    # 3. Final Fallback: Send message with public image link
    send_telegram_message(chat_id, f"🎨 Image generated: {public_url}")

def call_kora_chat(prompt: str, model: str = "kora-chat-0.5b") -> str:
    # 1. Try direct in-memory generation if running within gateway process
    try:
        import sys
        gw = sys.modules.get("gateway") or sys.modules.get("__main__")
        if gw and hasattr(gw, "CODER_MODEL") and hasattr(gw, "CODER_TOKENIZER"):
            tok = gw.CODER_TOKENIZER
            mdl = gw.CODER_MODEL
            is_coder = "coder" in model or "code" in prompt.lower()
            system_prompt = "You are Kora Coder, an expert AI programming assistant built by Kontyra." if is_coder else "You are Kora, an intelligent, helpful, and friendly AI assistant built by Kontyra."
            messages = [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": prompt}
            ]
            prompt_str = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
            device = next(mdl.parameters()).device
            inputs = tok([prompt_str], return_tensors="pt").to(device)
            import torch
            with torch.no_grad():
                out = mdl.generate(
                    **inputs,
                    max_new_tokens=300,
                    temperature=0.3 if is_coder else 0.7,
                    pad_token_id=tok.eos_token_id
                )
            reply = tok.decode(out[0][inputs.input_ids.shape[1]:], skip_special_tokens=True).strip()
            if reply:
                return reply
    except Exception as err:
        print(f"ℹ️ In-memory inference notice: {err}", flush=True)

    # 2. Try HTTP endpoint with Authorization headers
    api_key = os.environ.get("KORA_API_KEY", "")
    headers = {"Content-Type": "application/json"}
    if api_key:
        headers["Authorization"] = f"Bearer {api_key}"

    endpoints_to_try = [KORA_ENDPOINT, "http://127.0.0.1:7860", "https://professorceo-kora-ai.hf.space"]
    for ep in endpoints_to_try:
        url = f"{ep}/v1/chat/completions"
        try:
            req_data = json.dumps({
                "model": model,
                "messages": [{"role": "user", "content": prompt}],
                "max_tokens": 256
            }).encode("utf-8")
            req = urllib.request.Request(url, data=req_data, headers=headers)
            with urllib.request.urlopen(req, timeout=10) as resp:
                data = json.loads(resp.read().decode("utf-8"))
                return data["choices"][0]["message"]["content"]
        except Exception:
            continue
    return "Hello! I am Kora AI. How can I assist you today?"

def call_kora_image(prompt: str) -> dict:
    # 1. Try direct in-memory local PyTorch Stable Diffusion generation
    try:
        from vision.image_generator import generate_image
        res = generate_image(prompt)
        if res and res.get("success"):
            return res
    except Exception as err:
        print(f"ℹ️ In-memory image generation notice: {err}", flush=True)

    # 2. Try HTTP endpoint
    url = f"{KORA_ENDPOINT}/v1/images/generations"
    try:
        req_data = json.dumps({"prompt": prompt, "size": "1024x1024"}).encode("utf-8")
        req = urllib.request.Request(url, data=req_data, headers={"Content-Type": "application/json"})
        with urllib.request.urlopen(req, timeout=30) as resp:
            data = json.loads(resp.read().decode("utf-8"))
            return {"direct_url": data["data"][0]["url"]}
    except Exception:
        pass
    return {}

def handle_update(update: dict, token: str = None):
    global BOT_TOKEN
    if token:
        BOT_TOKEN = token

    message = update.get("message") or update.get("edited_message", {})
    text = message.get("text", "").strip()
    chat_id = message.get("chat", {}).get("id")

    if not chat_id or not text:
        return

    print(f"📩 Telegram Update [{chat_id}]: '{text}'", flush=True)

    if text in ["/start", "/help"]:
        send_telegram_message(
            chat_id,
            "🧠 Welcome to Kora AI on Telegram! (Powered by Kontyra)\n\n"
            "I am your autonomous AI assistant. Here is what I can do:\n"
            "• Send any question to chat with Kora Chat.\n"
            "• Type /code <prompt> for code generation.\n"
            "• Type /image <prompt> to generate 1024x1024 watermark-free images."
        )
    elif text.startswith("/image") or text.lower().startswith("create image"):
        prompt = text.replace("/image", "").replace("create image", "").strip()
        if not prompt:
            send_telegram_message(
                chat_id,
                "🎨 Please specify what image you would like to generate!\n\n"
                "Examples:\n"
                "• /image A sleek futuristic AI logo with glowing purple lights\n"
                "• /image A serene Japanese garden at sunrise, cinematic lighting"
            )
            return
        send_telegram_message(chat_id, f"🎨 Generating zero-watermark image for: {prompt}...")
        img_res = call_kora_image(prompt)
        file_path = img_res.get("file_path", "")
        direct_url = img_res.get("direct_url", "")
        photo_target = file_path if (file_path and os.path.exists(file_path)) else direct_url
        send_telegram_photo(chat_id, photo_path_or_url=photo_target, caption=f"✨ {prompt}")
    elif text.startswith("/code"):
        prompt = text.replace("/code", "").strip() or "Write a Python function to demonstrate binary search"
        send_telegram_message(chat_id, "💻 Generating code with Kora Coder...")
        reply = call_kora_chat(prompt, model="kora-coder-0.5b")
        send_telegram_message(chat_id, f"```\n{reply}\n```")
    else:
        reply = call_kora_chat(text, model="kora-chat-0.5b")
        send_telegram_message(chat_id, reply)

def run_bot():
    if not BOT_TOKEN:
        print("ℹ️ TELEGRAM_BOT_TOKEN environment variable not set. Telegram worker idle.", flush=True)
        return

    print("=" * 70, flush=True)
    print("🤖 KORA TELEGRAM BOT WORKER STARTED", flush=True)
    print(f"📡 Target Kora Gateway: {KORA_ENDPOINT}", flush=True)
    print("=" * 70, flush=True)

    test_telegram_connectivity()

    offset = 0
    while True:
        try:
            res = send_telegram_request("getUpdates", {"offset": offset, "timeout": 20}, timeout_sec=25.0)
            if res and res.get("ok"):
                for result in res.get("result", []):
                    offset = max(offset, result["update_id"] + 1)
                    handle_update(result, BOT_TOKEN)
            else:
                time.sleep(5)
        except Exception as e:
            safe_err = redact_token(str(e))
            print(f"⚠️ Telegram polling warning: {safe_err}", flush=True)
            time.sleep(5)

if __name__ == "__main__":
    run_bot()