app.py

import os
import re
import asyncio
import torch
import time
import httpx
import uvicorn
from collections import defaultdict
from fastapi import FastAPI, Request, BackgroundTasks
from fastapi.responses import JSONResponse, HTMLResponse
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftConfig, PeftModel
from duckduckgo_search import DDGS
from os_commands import ErozOS

ADAPTER_PATH = os.environ.get("ADAPTER_PATH")
BASE_MODEL = os.environ.get("BASE_MODEL")
HF_TOKEN = os.environ.get("HF_TOKEN")
CHANDRA_API_KEY = os.environ.get("CHANDRA_API_KEY")
BOT_TOKEN = os.environ.get("BOT_TOKEN")
WORKER_URL = os.environ.get("WORKER_URL")
WORKER_SECRET = os.environ.get("WORKER_SECRET", "")

ALLOWED_USERS = [int(x.strip()) for x in os.environ.get("ALLOWED_USERS", "").split(",") if x.strip()]

processed_messages = defaultdict(float)
MESSAGE_CACHE_TTL = 600

chat_histories = defaultdict(list)
MAX_HISTORY = 10

TELEGRAM_MESSAGE_LIMIT = 4096

ALLOWED_HTML_TAGS = {"b", "i", "u", "s", "code", "pre", "a", "tg-spoiler", "blockquote", "tg-emoji", "span"}

os_core = ErozOS()

SEARCH_KEYWORDS = [
    "latest", "current", "news", "today", "recent", "update", "2024",
    "2025", "2026", "now", "who is", "what is happening",
]

BAD_RESPONSE_PATTERNS = [
    "currentuser",
    "access_token",
    "refresh_token",
    "hashed_password",
    "john.doe@example.com",
    "127.0.0.1",
    "fukaiaefnk",
    " jdh",
    "j dh",
    "tellme somet",
    "tellyou somet",
]

SYSTEM_PROMPT = """You are EROZ AI, a helpful Telegram assistant created by DR.ZERO (Kavish Nethara).

Your main job is helping with A/L Science: Biology, Chemistry, and Physics. You can also answer normal questions clearly.

Rules:
- Answer the user's actual message directly.
- For greetings, reply briefly and naturally.
- For science questions, explain step by step in simple English unless the user asks for another language.
- Use Telegram HTML only when helpful: <b>, <i>, <code>, <pre>, and <blockquote>.
- Do not invent fake users, accounts, APIs, JSON fields, tokens, passwords, or currentUser examples unless the user explicitly asks about programming data structures.
- Do not switch into Chinese or other languages unless the user asks.
- If you are unsure, say so and give the best useful explanation."""


# --- Async HTTP Client (global) ---
async_client = None


async def get_http_client():
    """Get or create the async HTTP client."""
    global async_client
    if async_client is None or async_client.is_closed:
        async_client = httpx.AsyncClient(verify=False, timeout=60.0)
    return async_client


# --- Sync HTTP client for startup only ---
sync_client = httpx.Client(verify=False, timeout=60.0)


# --- HTML Sanitization (fixed — proper tag-aware approach) ---

def sanitize_html(text):
    """
    Convert markdown-style formatting to Telegram HTML and strip disallowed tags.
    Fixed: uses proper tag-by-tag parsing instead of broken greedy regex.
    """
    if not text:
        return ""

    # Convert markdown bold/italic/code to HTML
    text = re.sub(r'\*\*(.+?)\*\*', r'<b>\1</b>', text)
    text = re.sub(r'__(.+?)__', r'<i>\1</i>', text)
    text = re.sub(r'`([^`\n]+)`', r'<code>\1</code>', text)

    # Strip disallowed HTML tags while keeping their content
    def clean_tag(match):
        full_tag = match.group(0)
        is_closing = match.group(1)  # "/" or ""
        tag_name = match.group(2).split()[0].lower()  # Get just the tag name, ignore attributes

        if tag_name in ALLOWED_HTML_TAGS:
            return full_tag
        # Not allowed — remove the tag but keep what's inside
        return ""

    # Match opening and closing tags separately
    text = re.sub(r'<(/?)(\w[^>]*?)>', clean_tag, text)

    # Fix common broken HTML: unclosed tags
    text = _fix_unclosed_tags(text)

    return text


def _fix_unclosed_tags(text):
    """Ensure all opened tags are properly closed for Telegram HTML."""
    tag_stack = []
    # Find all tags
    for match in re.finditer(r'<(/?)(\w+)[^>]*>', text):
        is_closing = match.group(1) == "/"
        tag_name = match.group(2).lower()
        if tag_name not in ALLOWED_HTML_TAGS:
            continue
        if is_closing:
            # Remove from stack if present
            if tag_stack and tag_stack[-1] == tag_name:
                tag_stack.pop()
        else:
            tag_stack.append(tag_name)

    # Close any unclosed tags in reverse order
    for tag in reversed(tag_stack):
        text += f"</{tag}>"

    return text


def strip_html(text):
    """Remove all HTML tags, returning plain text."""
    return re.sub(r'<[^>]+>', '', text)


# --- Telegram API (async with worker fallback) ---

async def telegram_api(method, payload=None):
    """
    Send request via Netlify worker proxy.
    Falls back to direct Telegram API if worker fails.
    Passes WORKER_SECRET for authentication.
    """
    client = await get_http_client()
    data = {
        "path": method,
        "token": BOT_TOKEN,
        "body": payload or {},
    }
    # Pass secret if configured
    if WORKER_SECRET:
        data["secret"] = WORKER_SECRET

    for attempt in range(3):
        try:
            resp = await client.post(WORKER_URL, json=data)
            print(f"Telegram Worker {method}: {resp.status_code}")
            if resp.status_code == 200:
                return resp.json()
            # If worker returns auth error or server error, try fallback
            if resp.status_code in (403, 500, 502, 503):
                print(f"Worker returned {resp.status_code}, falling back to direct API")
                return await telegram_api_direct(method, payload)
            return resp.json()
        except Exception as e:
            print(f"Telegram Worker {method} attempt {attempt+1} FAILED: {e}")
            if attempt < 2:
                await asyncio.sleep(1)

    # All worker attempts failed — fallback to direct
    print(f"Worker {method}: ALL ATTEMPTS FAILED, using direct API")
    return await telegram_api_direct(method, payload)


async def telegram_api_direct(method, payload=None):
    """Direct call to Telegram API (fallback when worker is down)."""
    client = await get_http_client()
    url = f"https://api.telegram.org/bot{BOT_TOKEN}/{method}"
    try:
        resp = await client.post(url, json=payload or {})
        print(f"Telegram Direct {method}: {resp.status_code}")
        return resp.json()
    except Exception as e:
        print(f"Telegram Direct {method} FAILED: {e}")
        return {"ok": False, "description": str(e)}


def telegram_api_sync(method, payload=None):
    """Synchronous Telegram API call (for startup only)."""
    data = {
        "path": method,
        "token": BOT_TOKEN,
        "body": payload or {},
    }
    if WORKER_SECRET:
        data["secret"] = WORKER_SECRET
    try:
        resp = sync_client.post(WORKER_URL, json=data)
        print(f"Telegram Sync {method}: {resp.status_code}")
        return resp.json()
    except Exception as e:
        print(f"Telegram Sync {method} FAILED: {e}")
        return {"ok": False}


# --- Message Sending (async with HTML fallback) ---

async def send_message(chat_id, text, reply_to_message_id=None):
    """
    Send a message to Telegram.
    - Sanitizes HTML
    - Splits long messages
    - Falls back to plain text if HTML is rejected by Telegram
    """
    clean_text = sanitize_html(text)

    if len(clean_text) <= TELEGRAM_MESSAGE_LIMIT:
        await _send_single(chat_id, clean_text, reply_to_message_id)
    else:
        chunks = split_message(clean_text)
        for i, chunk in enumerate(chunks):
            rid = reply_to_message_id if i == 0 else None
            await _send_single(chat_id, chunk, rid)


async def _send_single(chat_id, text, reply_to_message_id=None):
    """Send a single message with HTML fallback."""
    payload = {"chat_id": chat_id, "text": text, "parse_mode": "HTML"}
    if reply_to_message_id:
        payload["reply_to_message_id"] = reply_to_message_id

    result = await telegram_api("sendMessage", payload)

    # If Telegram rejected our HTML, retry as plain text
    if not result.get("ok") and "can't parse" in result.get("description", "").lower():
        print(f"HTML rejected by Telegram, retrying as plain text")
        plain_text = strip_html(text)
        payload["text"] = plain_text
        payload.pop("parse_mode", None)
        result = await telegram_api("sendMessage", payload)

    return result


def split_message(text, max_len=TELEGRAM_MESSAGE_LIMIT):
    """Split long messages at newline boundaries, preserving HTML tag integrity."""
    if len(text) <= max_len:
        return [text]

    chunks = []
    lines = text.split("\n")
    current = ""

    for line in lines:
        if len(current) + len(line) + 1 > max_len:
            if current:
                chunks.append(current)
            # If a single line is too long, force-split it
            if len(line) > max_len:
                while len(line) > max_len:
                    chunks.append(line[:max_len])
                    line = line[max_len:]
                current = line
            else:
                current = line
        else:
            current = current + "\n" + line if current else line

    if current:
        chunks.append(current)

    return chunks if chunks else [text[:max_len]]


# --- Model Loading ---

def resolve_runtime_base_model(adapter_base_model):
    """Map training-only quantized bases to CPU-loadable runtime bases."""
    requested_base = BASE_MODEL
    force_base = os.environ.get("FORCE_BASE_MODEL") == "1"

    if requested_base and force_base:
        return requested_base

    if requested_base and requested_base.endswith("-Instruct"):
        return requested_base

    if adapter_base_model:
        normalized = adapter_base_model
        if normalized.startswith("unsloth/"):
            normalized = normalized.replace("unsloth/", "Qwen/", 1)
        if normalized.endswith("-bnb-4bit"):
            normalized = normalized[:-len("-bnb-4bit")]
        if "Qwen2.5-1.5B-Instruct" in normalized:
            return "Qwen/Qwen2.5-1.5B-Instruct"
        return normalized

    return requested_base


def load_model_hf():
    if not ADAPTER_PATH:
        raise RuntimeError("ADAPTER_PATH is not set. Set it to your LoRA repo, e.g. drxzero/eroz-ai-qwen2.5-1.5b-lora.")

    adapter_config = PeftConfig.from_pretrained(ADAPTER_PATH, token=HF_TOKEN)
    adapter_base_model = getattr(adapter_config, "base_model_name_or_path", None)
    base_model_name = resolve_runtime_base_model(adapter_base_model)

    if not base_model_name:
        raise RuntimeError("BASE_MODEL is not set and the adapter config does not declare a base_model_name_or_path.")

    if BASE_MODEL and adapter_base_model and BASE_MODEL != base_model_name:
        print(
            f"WARNING: BASE_MODEL={BASE_MODEL} does not match adapter base "
            f"{adapter_base_model}. Using runtime base {base_model_name}. "
            "Set FORCE_BASE_MODEL=1 only if you really want to override this.",
            flush=True,
        )

    print(f"Loading base model: {base_model_name}", flush=True)

    tokenizer = AutoTokenizer.from_pretrained(ADAPTER_PATH, trust_remote_code=True, token=HF_TOKEN)
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token = tokenizer.eos_token
    base_model = AutoModelForCausalLM.from_pretrained(
        base_model_name,
        device_map="cpu",
        dtype=torch.float32,
        trust_remote_code=True,
        token=HF_TOKEN,
    )
    model = PeftModel.from_pretrained(base_model, ADAPTER_PATH, token=HF_TOKEN)
    model.eval()
    return model, tokenizer


print("Loading model...", flush=True)
model, tokenizer = load_model_hf()
print("Model loaded!", flush=True)

# Test worker connection at startup (sync)
test = telegram_api_sync("getMe")
print(f"Worker test: {test}", flush=True)


# --- Web Search ---

def search_web(query, max_results=3):
    try:
        with DDGS() as ddgs:
            results = list(ddgs.text(query, max_results=max_results))
            if results:
                context = "\n".join([f"- {r['title']}: {r['body']}" for r in results])
                return context
    except Exception as e:
        print(f"Search error: {e}")
    return None


# --- Photo Download (async, passes secret) ---

async def download_telegram_file(file_id, min_size=1):
    """Download a Telegram file via worker proxy. Passes secret for auth."""
    try:
        file_info = await telegram_api("getFile", {"file_id": file_id})
        print(f"File info: {file_info}")
        if file_info.get("ok"):
            file_path = file_info["result"]["file_path"]
            # Build URL with secret param
            url = f"{WORKER_URL}?token={BOT_TOKEN}&file_path={file_path}"
            if WORKER_SECRET:
                url += f"&secret={WORKER_SECRET}"
            print(f"Downloading from Worker: {url}")
            client = await get_http_client()
            resp = await client.get(url)
            print(f"Worker download status: {resp.status_code}, size: {len(resp.content)}, content_type: {resp.headers.get('content-type', 'unknown')}")
            if resp.status_code == 200 and len(resp.content) >= min_size:
                return resp.content

            # Fallback: try direct Telegram download
            print(f"Worker download failed ({len(resp.content)} bytes), trying direct...")
            direct_url = f"https://api.telegram.org/file/bot{BOT_TOKEN}/{file_path}"
            resp = await client.get(direct_url)
            if resp.status_code == 200 and len(resp.content) >= min_size:
                return resp.content
            print(f"Direct download also failed: {resp.status_code}")
    except Exception as e:
        print(f"Telegram file download error: {e}")
    return None


async def download_telegram_photo(file_id):
    return await download_telegram_file(file_id, min_size=100)


# --- OCR with Chandra ---

async def ocr_with_chandra(image_bytes):
    try:
        client = await get_http_client()
        resp = await client.post(
            "https://www.datalab.to/api/v1/convert",
            headers={"X-API-Key": CHANDRA_API_KEY},
            files={"file": ("photo.jpg", image_bytes, "image/jpeg")},
            data={"output_format": "markdown"},
        )
        print(f"Chandra submit status: {resp.status_code}")
        if resp.status_code == 200:
            result = resp.json()
            check_url = result.get("request_check_url")
            print(f"Chandra check URL: {check_url}")

            if check_url:
                for i in range(30):
                    await asyncio.sleep(2)
                    poll_resp = await client.get(check_url, headers={"X-API-Key": CHANDRA_API_KEY})
                    poll_data = poll_resp.json()
                    status = poll_data.get("status")
                    print(f"Chandra poll {i+1}: status={status}")

                    if status == "complete":
                        md = poll_data.get("markdown") or ""
                        txt = poll_data.get("text") or ""
                        print(f"Chandra markdown length: {len(md)}")
                        print(f"Chandra text length: {len(txt)}")
                        print(f"Chandra result keys: {list(poll_data.keys())}")
                        if md:
                            return md
                        elif txt:
                            return txt
                        else:
                            print(f"Chandra full response: {poll_data}")
                            return str(poll_data)
                    elif status == "failed":
                        print(f"Chandra failed: {poll_data.get('error')}")
                        return None
        else:
            print(f"Chandra error: {resp.text}")
    except Exception as e:
        print(f"Chandra OCR error: {e}")
    return None


# --- Translation ---

async def translate_to_english(text):
    try:
        client = await get_http_client()
        resp = await client.get(
            "https://api.mymemory.translated.net/get",
            params={"q": text[:500], "langpair": "si|en"},
        )
        if resp.status_code == 200:
            result = resp.json()
            translated = result.get("responseData", {}).get("translatedText", "")
            print(f"Translation input length: {len(text)}")
            print(f"Translation output length: {len(translated)}")
            print(f"Translation sample: {translated[:200]}")
            if translated and translated.lower() != text.lower():
                return translated
    except Exception as e:
        print(f"Translation error: {e}")
    return text


# --- AI Response Generation ---

def response_looks_broken(response):
    """Catch common dataset/template leaks before they reach Telegram."""
    if not response:
        return True
    low = response.lower()
    if any(pattern in low for pattern in BAD_RESPONSE_PATTERNS):
        return True
    if low.count("tell me something") >= 2:
        return True
    if low.count("what did you") >= 2:
        return True
    if low.count("=") >= 5 and (low.count("'") + low.count('"')) >= 10:
        return True
    if response.count("\n") > 25 and len(response) > 1200:
        return True
    cjk_chars = sum(1 for ch in response if "\u4e00" <= ch <= "\u9fff")
    if len(response) > 120 and cjk_chars / max(len(response), 1) > 0.15:
        return True
    private_chars = sum(1 for ch in response if "\ue000" <= ch <= "\uf8ff")
    if private_chars:
        return True
    fullwidth_chars = sum(1 for ch in response if "\uff00" <= ch <= "\uffef")
    if fullwidth_chars > 3:
        return True
    return False


def clean_model_response(response):
    if not response:
        return ""
    response = response.replace("<|im_end|>", "").replace("<|endoftext|>", "")
    response = response.strip()

    # Some generations keep talking as the next user/system turn. Cut that off.
    stop_markers = ["\nUser:", "\nuser:", "\nAssistant:", "\nassistant:", "\nSystem:", "\nsystem:"]
    for marker in stop_markers:
        idx = response.find(marker)
        if idx > 0:
            response = response[:idx].strip()
    return response


def build_messages(user_input, chat_id, reply_context=None):
    system_content = SYSTEM_PROMPT

    if reply_context:
        system_content += (
            "\n\nThe user is replying to this earlier Telegram message. "
            "Use it only as context and answer the new question:\n"
            f"<blockquote>{reply_context[:1200]}</blockquote>"
        )

    needs_search = any(kw in user_input.lower() for kw in SEARCH_KEYWORDS)
    if needs_search:
        search_results = search_web(user_input)
        if search_results:
            system_content += (
                "\n\nRecent web context:\n"
                f"{search_results}\n\nUse this only if it is relevant."
            )

    messages = [{"role": "system", "content": system_content}]

    # Keep only clean history. Bad generations can otherwise teach the next turn
    # to repeat the same broken dataset artifact.
    for item in chat_histories.get(chat_id, [])[-8:]:
        content = item.get("content", "")
        if item.get("role") == "assistant" and response_looks_broken(content):
            continue
        messages.append({"role": item.get("role", "user"), "content": content[:1500]})

    messages.append({"role": "user", "content": user_input})
    return messages


def fallback_response(user_input):
    low = user_input.lower().strip()
    if low in {"hi", "hello", "hey", "helo"}:
        return "Hello! I am <b>EROZ AI</b>. Ask me any Biology, Chemistry, Physics, or general question."
    if "cell" in low:
        return (
            "A <b>cell</b> is the basic structural and functional unit of life. "
            "All living organisms are made of one or more cells.\n\n"
            "<b>Main parts:</b>\n"
            "- <b>Cell membrane</b>: controls what enters and leaves the cell.\n"
            "- <b>Cytoplasm</b>: jelly-like material where many reactions happen.\n"
            "- <b>Nucleus</b>: contains DNA and controls cell activities.\n"
            "- <b>Mitochondria</b>: release energy by respiration.\n\n"
            "Plant cells also have a <b>cell wall</b>, <b>chloroplasts</b>, and a large <b>vacuole</b>."
        )
    return (
        "I could not generate a clean answer for that message. "
        "Please ask again with a little more detail, or use <code>/clear</code> and retry."
    )


def generate_from_text(prompt_text, max_new_tokens=384):
    inputs = tokenizer(prompt_text, return_tensors="pt").to(model.device)
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=False,
            repetition_penalty=1.18,
            no_repeat_ngram_size=4,
            eos_token_id=tokenizer.eos_token_id,
            pad_token_id=tokenizer.pad_token_id,
        )
    response = tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True)
    return clean_model_response(response)


def build_plain_prompt(user_input, reply_context=None):
    context = ""
    if reply_context:
        context = f"Context: {reply_context[:1000]}\n\n"
    return (
        "You are EROZ AI, a helpful A/L Science assistant. "
        "Answer only the user's question in clear English. "
        "Do not list paraphrases, fake variables, tokens, or training examples.\n\n"
        f"{context}Question: {user_input}\nAnswer:"
    )


def generate_ai_response(user_input, chat_id, reply_context=None):
    messages = build_messages(user_input, chat_id, reply_context)
    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    response = generate_from_text(text, max_new_tokens=512)

    if response_looks_broken(response):
        print(f"Broken chat-template response detected, retrying plain prompt: {response[:200]}")
        response = generate_from_text(build_plain_prompt(user_input, reply_context), max_new_tokens=384)

    if response_looks_broken(response):
        print(f"Broken plain-prompt response detected, using fallback: {response[:200]}")
        response = fallback_response(user_input)

    chat_histories[chat_id].append({"role": "user", "content": user_input})
    chat_histories[chat_id].append({"role": "assistant", "content": response})

    if len(chat_histories[chat_id]) > MAX_HISTORY * 2:
        chat_histories[chat_id] = chat_histories[chat_id][-MAX_HISTORY * 2:]

    return response


# --- FastAPI App ---

app = FastAPI()


@app.on_event("shutdown")
async def shutdown():
    """Close the async HTTP client on shutdown."""
    global async_client
    if async_client and not async_client.is_closed:
        await async_client.aclose()


@app.post("/webhook")
async def webhook(request: Request, background_tasks: BackgroundTasks):
    try:
        data = await request.json()
        message = data.get("message", {})
        chat_id = message.get("chat", {}).get("id")
        text = message.get("text", "").strip()
        message_id = message.get("message_id")
        photos = message.get("photo")
        document = message.get("document")
        caption = message.get("caption", "").strip()  # Fixed: capture photo captions

        reply_to = message.get("reply_to_message")
        reply_context = None
        if reply_to:
            reply_context = reply_to.get("text", "")

        print(f"Webhook: chat_id={chat_id}, text={text}, caption={caption}, msg_id={message_id}, has_photo={bool(photos)}, has_document={bool(document)}, reply_to={bool(reply_to)}")
        if reply_context:
            print(f"Reply context: {reply_context[:200]}")

        if not chat_id:
            return JSONResponse({"ok": True})

        if chat_id not in ALLOWED_USERS:
            await send_message(chat_id, "Access denied. This bot is private.")
            return JSONResponse({"ok": True})

        # Dedup check
        msg_key = f"{chat_id}_{message_id}"
        now = time.time()

        if msg_key in processed_messages and (now - processed_messages[msg_key]) < MESSAGE_CACHE_TTL:
            print(f"Duplicate message skipped: {msg_key}")
            return JSONResponse({"ok": True})

        processed_messages[msg_key] = now

        # Clean old cache entries
        for key in list(processed_messages.keys()):
            if now - processed_messages[key] > MESSAGE_CACHE_TTL:
                del processed_messages[key]

        # --- Fire pending reminders ---
        fired_reminders = os_core.check_reminders(chat_id)
        for r in fired_reminders:
            await send_message(chat_id, f"🔔 <b>Reminder:</b> {r['msg']}")

        # --- Background Task Helpers ---
        async def process_photo_task(chat_id, photos, caption):
            try:
                largest_photo = max(photos, key=lambda p: p.get("file_size", 0))
                file_id = largest_photo.get("file_id")
                if file_id:
                    image_bytes = await download_telegram_photo(file_id)
                    if image_bytes:
                        ocr_text = await ocr_with_chandra(image_bytes)
                        if ocr_text:
                            english_text = await translate_to_english(ocr_text)
                            if caption:
                                prompt = f"User's question: {caption}\n\nExtracted text from image (translated to English):\n\n{english_text}\n\nOriginal Sinhala text:\n\n{ocr_text}\n\nAnswer the user's question based on this image content."
                            else:
                                prompt = f"Extracted text from image (translated to English):\n\n{english_text}\n\nOriginal Sinhala text:\n\n{ocr_text}\n\nAnswer the question from this image. If it's a multiple choice question, explain which answer is correct and why. Explain in simple English."
                            
                            response = await asyncio.to_thread(generate_ai_response, prompt, chat_id)
                            await send_message(chat_id, response)
                        else:
                            await send_message(chat_id, "Could not read the image. Please try again.")
                    else:
                        await send_message(chat_id, "Could not download the photo. Please try again.")
            except Exception as e:
                print(f"AI generation error (photo task): {e}")
                await send_message(chat_id, f"<b>Error:</b> Could not generate response. Please try again.\n<code>{str(e)[:200]}</code>")

        async def process_document_task(chat_id, document, caption):
            try:
                file_id = document.get("file_id")
                filename = document.get("file_name") or f"telegram_{file_id}.bin"
                size = document.get("file_size", 0)
                if not file_id:
                    await send_message(chat_id, "<b>Error:</b> Telegram did not include a file id.")
                    return

                if size and size > 45 * 1024 * 1024:
                    await send_message(chat_id, "<b>Error:</b> File is too large. Please keep uploads under 45 MB.")
                    return

                content = await download_telegram_file(file_id, min_size=1)
                if not content:
                    await send_message(chat_id, "<b>Error:</b> Could not download the document.")
                    return

                saved = os_core.save_uploaded_file(filename, content)
                await send_message(
                    chat_id,
                    f"<b>File saved:</b> <code>{saved['name']}</code>\n"
                    f"<i>Size:</i> <code>{saved['size']} bytes</code>\n"
                    f"<i>SHA256:</i> <code>{saved['sha256']}</code>",
                )

                if caption.lower().startswith("/hf upload"):
                    extra = caption[len("/hf upload"):].strip()
                    command = f"/hf upload {saved['name']}"
                    if extra:
                        command = f"{command} {extra}"
                    hf_response = await asyncio.to_thread(os_core.handle, command, chat_id)
                    if hf_response:
                        await send_message(chat_id, hf_response)
            except Exception as e:
                print(f"Document task error: {e}")
                await send_message(chat_id, f"<b>Error:</b> Could not process document.\n<code>{str(e)[:200]}</code>")

        async def process_text_task(chat_id, text, reply_context):
            try:
                response = await asyncio.to_thread(generate_ai_response, text, chat_id, reply_context)
                await send_message(chat_id, response)
            except Exception as e:
                print(f"AI generation error (text task): {e}")
                await send_message(chat_id, f"<b>Error:</b> Could not generate response. Please try again.\n<code>{str(e)[:200]}</code>")

        # --- Handle message ---
        if text == "/start":
            chat_histories[chat_id] = []
            await send_message(chat_id, "<b>Hello! I am EROZ AI</b>, your A/L Science assistant.\n\nI can help with <b>Biology</b>, <b>Chemistry</b>, and <b>Physics</b>.\n\nCreated by <b>DR.ZERO</b> (Kavish Nethara).\n\nSend me any question or photo to begin!")
        elif text.lower() in ["/clear", "clear", "reset"]:
            chat_histories[chat_id] = []
            await send_message(chat_id, "<b>Chat history cleared!</b>")
        elif photos:
            await send_message(chat_id, "<i>Processing your photo... (this may take a minute)</i>")
            background_tasks.add_task(process_photo_task, chat_id, photos, caption)
        elif document:
            await send_message(chat_id, "<i>Saving your file...</i>")
            background_tasks.add_task(process_document_task, chat_id, document, caption)
        elif os_core.is_command(text):
            os_response = os_core.handle(text, chat_id, generate_ai_fn=generate_ai_response, search_web_fn=search_web)
            if os_response:
                await send_message(chat_id, os_response)
            else:
                background_tasks.add_task(process_text_task, chat_id, text, reply_context)
        else:
            background_tasks.add_task(process_text_task, chat_id, text, reply_context)

        return JSONResponse({"ok": True})
    except Exception as e:
        print(f"Webhook error: {e}")
        # Try to notify the user even on critical errors
        try:
            if chat_id:
                await send_message(chat_id, "<b>Something went wrong.</b> Please try again.")
        except Exception:
            pass
        return JSONResponse({"ok": False, "error": str(e)})


@app.get("/")
async def home():
    return HTMLResponse("<h1>Eroz AI Bot - Running</h1>")


@app.get("/health")
async def health():
    return JSONResponse({"status": "ok"})


print("Bot ready! Webhook: https://drxzero-eroz-ai-1-5b.hf.space/webhook", flush=True)

uvicorn.run(app, host="0.0.0.0", port=7860)