from utils.ai_brain import ai_handle_query
import requests
from pydantic import BaseModel
from fastapi import FastAPI, BackgroundTasks
import threading

app = FastAPI()
class ChatRequest(BaseModel):
    message: str
    chat_id: str | None = None

import uvicorn
from dotenv import load_dotenv
load_dotenv()
import os

MODE = os.getenv("MODE", "full")  # default = full (for local dev)

from agents.openrouter_agent import fetch_openrouter_models
from agents.hf_agent import fetch_hf_trending_models
from memory.dataset_manager import (
    download_memory,
    upload_memory,
    is_new_model,
    add_model
)
from notifier.telegram import send_telegram_message
from utils.ranker import score_openrouter, score_huggingface
from utils.query_engine import (
    get_best_free_openrouter,
    get_best_hf_model,
    get_best_free_model,
    get_top_models_by_category,
    get_trending_models
)
from notifier.telegram_listener import send_message
from utils.auto_update import run_autonomous_update

# FAST API
@app.get("/update")
def trigger_update():
    print("⚡ External update triggered...")

    threading.Thread(
        target=run_update_cycle,
        daemon=True
    ).start()

    return {"status": "update started"}

@app.post("/chat")
def chat_endpoint(req: ChatRequest):
    print(f"💬 Incoming chat: {req.message}")

    # Fetch latest models (fresh context)
    or_models = fetch_openrouter_models()
    hf_models = fetch_hf_trending_models()

    try:
        reply = ai_handle_query(
            user_query=req.message,
            or_models=or_models,
            hf_models=hf_models,
            chat_id=req.chat_id
        )

        return {
            "reply": reply
        }

    except Exception as e:
        print("❌ Chat error:", str(e))
        return {
            "reply": "⚠️ Error processing request"
        }

# Deduplication cache for Telegram update IDs to prevent loops
PROCESSED_UPDATES = set()

def process_telegram_background(update_id: int, chat_id: int, text: str):
    try:
        print(f"⚙️ Processing Telegram query in background for update_id {update_id}: '{text}'")
        or_models = fetch_openrouter_models()
        hf_models = fetch_hf_trending_models()

        reply = ai_handle_query(
            user_query=text,
            or_models=or_models,
            hf_models=hf_models,
            chat_id=str(chat_id)
        )

        send_message(chat_id, reply)
        print(f"✅ Background processing done for update_id {update_id}")
    except Exception as e:
        print(f"❌ Background Telegram error for update_id {update_id}: {e}")

@app.post("/telegram")
async def telegram_webhook(update: dict, background_tasks: BackgroundTasks):
    try:
        update_id = update.get("update_id")
        message = update.get("message", {})
        chat_id = message.get("chat", {}).get("id")
        text = message.get("text", "")

        if not text or not chat_id or update_id is None:
            return {"status": "ignored"}

        # Deduplication check
        if update_id in PROCESSED_UPDATES:
            print(f"⏭️ Duplicate update_id {update_id} detected. Skipping processing.")
            return {"status": "already_processed"}

        PROCESSED_UPDATES.add(update_id)
        # Cap set size to prevent memory leaks
        if len(PROCESSED_UPDATES) > 500:
            PROCESSED_UPDATES.clear()
            PROCESSED_UPDATES.add(update_id)

        print(f"📩 Telegram webhook received update_id {update_id}: {text}")

        # Offload execution to a background task so we can return 200 OK immediately
        background_tasks.add_task(process_telegram_background, update_id, chat_id, text)

        return {"status": "ok"}

    except Exception as e:
        print(f"❌ Telegram webhook error: {e}")
        return {"status": "error"}

@app.get("/")
def home():
    return {"status": "AI Radar is running"}


def format_model_message(source, model_name, extra_info=None):
    msg = f"🚀 *New AI Model Detected!*\n\n"
    msg += f"*Name:* `{model_name}`\n"
    msg += f"*Source:* {source.capitalize()}\n"

    # --- FREE DETECTION ---
    is_free = "free" in model_name.lower()

    if extra_info:
        # Context
        context = extra_info.get("context")
        if context:
            msg += f"*Context:* {context} tokens\n"

        # Pricing
        pricing = extra_info.get("pricing", {})
        if pricing:
            msg += f"*Free:* {'✅ Yes' if is_free else '❌ Paid'}\n"

    # --- USE CASE DETECTION ---
    if "embed" in model_name:
        use_case = "Embeddings"
    elif "vision" in model_name or "clip" in model_name:
        use_case = "Vision"
    elif "chat" in model_name or "gpt" in model_name:
        use_case = "Chat / LLM"
    else:
        use_case = "General"

    msg += f"*Use Case:* {use_case}\n"
    
    if "score" in extra_info:
        stars = "⭐" * min(extra_info["score"], 5)
        msg += f"*Rank:* {stars}\n"

    msg += "\n🔥 *AI Radar Insight:*"
    if is_free:
        msg += "\nThis is a FREE model — worth testing immediately."
    else:
        msg += "\nPaid model — check performance before using."

    msg += "\n\n#AI #NewModel"
    return msg

def run_update_cycle():
    print("📦 Loading memory...")
    memory = download_memory()

    if len(memory) == 0:
        send_telegram_message("✅ Telegram is connected successfully!")

    print(f"🧠 Memory size: {len(memory)}")

    # Fetch models
    or_models = fetch_openrouter_models()
    hf_models = fetch_hf_trending_models()

    new_models = []

    # OpenRouter
    for m in or_models:
        if is_new_model(memory, m["name"]):
            score = score_openrouter(m)
            m["score"] = score
            new_models.append(("openrouter", m))
            memory = add_model(memory, m["name"], "openrouter")

    # HuggingFace
    for m in hf_models:
        if is_new_model(memory, m["name"]):
            score = score_huggingface(m)
            m["score"] = score
            new_models.append(("huggingface", m))
            memory = add_model(memory, m["name"], "huggingface")

    # Counts
    or_count = sum(1 for s, _ in new_models if s == "openrouter")
    hf_count = sum(1 for s, _ in new_models if s == "huggingface")

    print(f"\n📊 New OpenRouter models: {or_count}")
    print(f"📊 New Hugging Face models: {hf_count}")

    print(f"🆕 Total new models: {len(new_models)}")

    # Send alerts
    for source, model in new_models[:5]:
        message = format_model_message(source, model["name"], model)
        send_telegram_message(message)

    # Save memory
    upload_memory(memory)

    return or_models, hf_models

if __name__ == "__main__":

    if MODE == "update":
        print("⚡ Running UPDATE mode...")
        run_update_cycle()

    elif MODE == "bot" or MODE == "full":
        print(f"🚀 Starting Uvicorn server in {MODE} mode...")
        # The @app.on_event("startup") will safely handle starting the bot thread for us.
        uvicorn.run(app, host="0.0.0.0", port=7860)

@app.on_event("startup")
def startup_event():
    print("🚀 AI Radar is initializing...")

    # Define a small wrapper to run the startup sequence in the background
    def background_startup():
        try:
            print("📦 Background: Running initial update cycle...")
            run_update_cycle()
            
            print("🌐 Registering Telegram webhook...")
            bot_token = os.getenv("TELEGRAM_BOT_TOKEN")
            api_base = os.getenv("TELEGRAM_API_BASE", "https://api.telegram.org").rstrip("/")
            wh_url = "https://shinayush-ai-radar.hf.space/telegram"
            
            webhook_registered = False
            for attempt in range(1, 6):
                try:
                    print(f"📡 Webhook registration attempt {attempt}/5...")
                    resp = requests.get(
                        f"{api_base}/bot{bot_token}/setWebhook?url={wh_url}",
                        timeout=30
                    )
                    if resp.status_code == 200 and resp.json().get("ok"):
                        print("✅ Telegram webhook registered successfully")
                        webhook_registered = True
                        break
                    else:
                        print(f"⚠️ Webhook response error on attempt {attempt}: {resp.text}")
                except Exception as e:
                    print(f"⚠️ Webhook registration attempt {attempt} failed: {e}")
                
                if attempt < 5:
                    import time
                    time.sleep(5)
            
            if not webhook_registered:
                print("❌ Webhook registration failed completely after 5 attempts.")

            print("✅ Background: All systems active!")
        except Exception as e:
            print(f"❌ Background Startup Error: {e}")

    # Kick off the heavy lifting in a separate thread so the server starts NOW
    threading.Thread(target=background_startup, daemon=True).start()
    
    print("✅ Web Server is live! (Background tasks starting...)")