"""
MundiX API Backend v2.0
Central proxy for MundiX CLI: authentication, billing, sessions, AI proxy, collective memory.
Runs on FastAPI + PostgreSQL + Redis.

ONBOARDING FUNNEL:
  1. Install → CLI auto-registers device (machine_id, IP, OS) → gets temp api_key
  2. 3 FREE queries (anonymous, no signup)
  3. After 3 queries → MUST register (email) to continue
  4. Registration → +5 FREE queries (total 8)
  5. After 8 queries → MUST pay to continue (min R$ 49.90)
  6. Machine is blocked until payment

IMPORTANT: This backend is INDEPENDENT from MundiX Sentinel (enterprise).
MundiX CLI is for pentesters and students.
"""

import os
import time
import uuid
import hashlib
import secrets
import json
import platform
from datetime import datetime, timedelta
from contextlib import asynccontextmanager

from fastapi import FastAPI, HTTPException, Request, Depends, Header
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse, JSONResponse
from pydantic import BaseModel, Field
from typing import Optional, List

import asyncpg
import redis.asyncio as aioredis
import mercadopago
from openai import OpenAI

# ─── Config ──────────────────────────────────────────────────────────

HF_TOKEN = os.environ.get("HF_TOKEN", "")
DATABASE_URL = os.environ.get("DATABASE_URL", "postgresql://mundix:mundix_secret@localhost:5432/mundix_api")
REDIS_URL = os.environ.get("REDIS_URL", "redis://localhost:6379/1")
MP_ACCESS_TOKEN = os.environ.get("MP_ACCESS_TOKEN", "")
MP_PUBLIC_KEY = os.environ.get("MP_PUBLIC_KEY", "")

# Pricing (in credits: 1 credit ≈ 1 token)
FREE_ANONYMOUS_QUERIES = 3       # Before registration
FREE_REGISTERED_QUERIES = 5      # After registration (total = 3 + 5 = 8)
FREE_ANONYMOUS_CREDITS = 6000    # ~3 queries worth
FREE_REGISTERED_CREDITS = 10000  # ~5 queries worth (added on registration)
MIN_PURCHASE = 500000            # R$ 49.90 = 500k credits

CREDIT_PACKAGES = {
    "starter": {"credits": 500000, "price_brl": 49.90, "description": "~250 queries"},
    "pro": {"credits": 1200000, "price_brl": 99.90, "description": "~600 queries"},
    "enterprise": {"credits": 3000000, "price_brl": 199.90, "description": "~1500 queries"},
}

# Model registry (internal — maps codenames to real HF model IDs)
MODEL_REGISTRY = {
    "mx-redhawk-7b": {
        "hf_id": "DeepHat/DeepHat-V1-7B",
        "provider": "featherless-ai",
        "cost_multiplier": 1.0,
    },
    "mx-phantom-7b": {
        "hf_id": "WhiteRabbitNeo/WhiteRabbitNeo-V3-7B",
        "provider": "featherless-ai",
        "cost_multiplier": 1.0,
    },
    "mx-spectre-7b": {
        "hf_id": "WhiteRabbitNeo/WhiteRabbitNeo-2.5-Qwen-2.5-Coder-7B",
        "provider": "novita",
        "cost_multiplier": 0.8,
    },
    "mx-venom-8b": {
        "hf_id": "WhiteRabbitNeo/Llama-3-WhiteRabbitNeo-8B-v2.0",
        "provider": "novita",
        "cost_multiplier": 0.8,
    },
    "mx-cipher-7b": {
        "hf_id": "Qwen/Qwen2.5-Coder-7B-Instruct",
        "provider": "auto",
        "cost_multiplier": 0.5,
    },
}

FALLBACK_ORDER = ["mx-redhawk-7b", "mx-phantom-7b", "mx-spectre-7b", "mx-venom-8b", "mx-cipher-7b"]

# ─── Database Pool & Redis ───────────────────────────────────────────

db_pool: asyncpg.Pool = None
redis_client: aioredis.Redis = None
hf_client: OpenAI = None


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Startup/shutdown lifecycle."""
    global db_pool, redis_client, hf_client

    # Connect to PostgreSQL
    db_pool = await asyncpg.create_pool(DATABASE_URL, min_size=2, max_size=10)

    # Initialize database schema
    async with db_pool.acquire() as conn:
        await conn.execute(SCHEMA_SQL)

    # Connect to Redis
    redis_client = aioredis.from_url(REDIS_URL, decode_responses=True)

    # Initialize HuggingFace client
    hf_client = OpenAI(
        base_url="https://router.huggingface.co/v1",
        api_key=HF_TOKEN,
    )

    yield

    # Shutdown
    await db_pool.close()
    await redis_client.close()


app = FastAPI(
    title="MundiX API",
    description="Backend for MundiX CLI - AI-Powered Cybersecurity Copilot",
    version="2.0.0",
    lifespan=lifespan,
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# ─── Database Schema ─────────────────────────────────────────────────

SCHEMA_SQL = """
-- Devices table: tracks every installation
CREATE TABLE IF NOT EXISTS devices (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    machine_id VARCHAR(128) UNIQUE NOT NULL,
    ip_address VARCHAR(45),
    hostname VARCHAR(255),
    os_info VARCHAR(255),
    api_key VARCHAR(64) UNIQUE NOT NULL,
    api_key_hash VARCHAR(128) UNIQUE NOT NULL,
    user_id UUID,
    credits BIGINT DEFAULT 6000,
    total_queries BIGINT DEFAULT 0,
    total_tokens_used BIGINT DEFAULT 0,
    max_free_queries INT DEFAULT 3,
    status VARCHAR(20) DEFAULT 'anonymous',
    is_blocked BOOLEAN DEFAULT FALSE,
    created_at TIMESTAMP DEFAULT NOW(),
    last_active TIMESTAMP DEFAULT NOW()
);

-- Users table: registered users (email required)
CREATE TABLE IF NOT EXISTS users (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    email VARCHAR(255) UNIQUE NOT NULL,
    name VARCHAR(255),
    password_hash VARCHAR(255),
    credits BIGINT DEFAULT 0,
    total_queries BIGINT DEFAULT 0,
    total_tokens_used BIGINT DEFAULT 0,
    plan VARCHAR(20) DEFAULT 'free',
    is_active BOOLEAN DEFAULT TRUE,
    created_at TIMESTAMP DEFAULT NOW(),
    last_active TIMESTAMP DEFAULT NOW()
);

-- Sessions table
CREATE TABLE IF NOT EXISTS sessions (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    device_id UUID REFERENCES devices(id),
    user_id UUID,
    project_name VARCHAR(255),
    model VARCHAR(64) DEFAULT 'mx-redhawk-7b',
    messages JSONB DEFAULT '[]'::jsonb,
    context JSONB DEFAULT '{}'::jsonb,
    created_at TIMESTAMP DEFAULT NOW(),
    updated_at TIMESTAMP DEFAULT NOW()
);

-- Usage log
CREATE TABLE IF NOT EXISTS usage_log (
    id BIGSERIAL PRIMARY KEY,
    device_id UUID REFERENCES devices(id),
    user_id UUID,
    session_id UUID,
    model VARCHAR(64),
    input_tokens INT,
    output_tokens INT,
    credits_used INT,
    query_text TEXT,
    created_at TIMESTAMP DEFAULT NOW()
);

-- Collective memory
CREATE TABLE IF NOT EXISTS collective_memory (
    id BIGSERIAL PRIMARY KEY,
    device_id UUID,
    user_id UUID,
    category VARCHAR(64),
    query TEXT NOT NULL,
    response TEXT NOT NULL,
    commands JSONB DEFAULT '[]'::jsonb,
    success BOOLEAN DEFAULT TRUE,
    upvotes INT DEFAULT 0,
    created_at TIMESTAMP DEFAULT NOW()
);

-- Device activation codes (for linking device to user account)
CREATE TABLE IF NOT EXISTS device_activations (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    device_code VARCHAR(8) UNIQUE NOT NULL,
    machine_id VARCHAR(128),
    device_id UUID,
    user_id UUID,
    api_key VARCHAR(64),
    status VARCHAR(20) DEFAULT 'pending',
    created_at TIMESTAMP DEFAULT NOW(),
    expires_at TIMESTAMP DEFAULT NOW() + INTERVAL '15 minutes'
);

-- Indexes
CREATE INDEX IF NOT EXISTS idx_devices_machine_id ON devices(machine_id);
CREATE INDEX IF NOT EXISTS idx_devices_api_key_hash ON devices(api_key_hash);
CREATE INDEX IF NOT EXISTS idx_devices_user_id ON devices(user_id);
CREATE INDEX IF NOT EXISTS idx_users_email ON users(email);
CREATE INDEX IF NOT EXISTS idx_sessions_device_id ON sessions(device_id);
CREATE INDEX IF NOT EXISTS idx_usage_log_device_id ON usage_log(device_id);
CREATE INDEX IF NOT EXISTS idx_collective_memory_category ON collective_memory(category);
CREATE INDEX IF NOT EXISTS idx_device_activations_code ON device_activations(device_code);

-- Payments table: Mercado Pago integration
CREATE TABLE IF NOT EXISTS payments (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    device_id UUID REFERENCES devices(id),
    user_id UUID,
    mp_preference_id VARCHAR(100),
    mp_payment_id VARCHAR(100),
    package VARCHAR(20) NOT NULL,
    credits INTEGER NOT NULL,
    amount DECIMAL(10,2) NOT NULL,
    status VARCHAR(20) DEFAULT 'pending',
    mp_status VARCHAR(50),
    mp_detail JSONB,
    created_at TIMESTAMP DEFAULT NOW(),
    updated_at TIMESTAMP DEFAULT NOW()
);
CREATE INDEX IF NOT EXISTS idx_payments_device ON payments(device_id);
CREATE INDEX IF NOT EXISTS idx_payments_mp_preference ON payments(mp_preference_id);
CREATE INDEX IF NOT EXISTS idx_payments_mp_payment ON payments(mp_payment_id);
CREATE INDEX IF NOT EXISTS idx_payments_status ON payments(status);
"""


# ─── Pydantic Models ─────────────────────────────────────────────────

class DeviceRegisterRequest(BaseModel):
    """Auto-registration on first install."""
    machine_id: str
    ip_address: Optional[str] = None
    hostname: Optional[str] = None
    os_info: Optional[str] = None

class DeviceLinkRequest(BaseModel):
    """Link device to user account (after registration on website)."""
    device_code: str
    machine_id: str

class UserRegisterRequest(BaseModel):
    """Register a new user account (from website)."""
    email: str
    name: Optional[str] = None
    device_code: Optional[str] = None

class DeviceActivationComplete(BaseModel):
    """Complete device activation from website."""
    device_code: str
    email: str
    name: Optional[str] = None

class ChatRequest(BaseModel):
    model: str = "mx-redhawk-7b"
    messages: List[dict]
    session_id: Optional[str] = None
    project_name: Optional[str] = None
    stream: bool = True
    max_tokens: int = 4096
    temperature: float = 0.7

class MemoryEntry(BaseModel):
    category: str
    query: str
    response: str
    commands: List[str] = []
    success: bool = True

class PixPaymentRequest(BaseModel):
    """Request to create a direct PIX payment."""
    package: str
    payer_email: Optional[str] = None
    payer_first_name: Optional[str] = None
    payer_last_name: Optional[str] = None
    payer_cpf: Optional[str] = None



# ─── Auth Helpers ─────────────────────────────────────────────────────

def hash_api_key(api_key: str) -> str:
    return hashlib.sha256(api_key.encode()).hexdigest()

def generate_api_key() -> str:
    return f"mx_{secrets.token_hex(24)}"

def generate_device_code() -> str:
    return secrets.token_hex(4).upper()


async def get_device_by_api_key(api_key: str) -> dict:
    """Look up device by API key (uses Redis cache)."""
    key_hash = hash_api_key(api_key)

    # Check Redis cache first
    cached = await redis_client.get(f"device:{key_hash}")
    if cached:
        return json.loads(cached)

    # Query database
    async with db_pool.acquire() as conn:
        row = await conn.fetchrow(
            "SELECT * FROM devices WHERE api_key_hash = $1",
            key_hash,
        )
        if row:
            device = dict(row)
            device["id"] = str(device["id"])
            device["user_id"] = str(device["user_id"]) if device["user_id"] else None
            device["created_at"] = device["created_at"].isoformat()
            device["last_active"] = device["last_active"].isoformat()
            # Cache for 2 minutes
            await redis_client.setex(f"device:{key_hash}", 120, json.dumps(device))
            return device
    return None


async def get_current_device(authorization: str = Header(None)) -> dict:
    """Dependency to get the current authenticated device."""
    if not authorization:
        raise HTTPException(status_code=401, detail="Missing Authorization header")

    api_key = authorization.replace("Bearer ", "").strip()

    if not api_key.startswith("mx_"):
        raise HTTPException(status_code=401, detail="Invalid API key format")

    device = await get_device_by_api_key(api_key)
    if not device:
        raise HTTPException(status_code=401, detail="Invalid or unknown device")

    if device.get("is_blocked"):
        raise HTTPException(
            status_code=403,
            detail="Device blocked. Please purchase credits at https://chat.mundix.dev/billing"
        )

    return device



async def get_current_device_allow_blocked(authorization: str = Header(None)) -> dict:
    """Dependency to get the current authenticated device, even if blocked.
    Used for billing endpoints so blocked users can still purchase credits."""
    if not authorization:
        raise HTTPException(status_code=401, detail="Missing Authorization header")
    api_key = authorization.replace("Bearer ", "").strip()
    if not api_key.startswith("mx_"):
        raise HTTPException(status_code=401, detail="Invalid API key format")
    device = await get_device_by_api_key(api_key)
    if not device:
        raise HTTPException(status_code=401, detail="Invalid or unknown device")
    # NOTE: We do NOT check is_blocked here — blocked users must be able to buy credits
    return device

# ═══════════════════════════════════════════════════════════════════════
# STEP 1: AUTO-REGISTER DEVICE (zero friction — on first install)
# ═══════════════════════════════════════════════════════════════════════

@app.post("/auth/device/register")
async def register_device(req: DeviceRegisterRequest):
    """Auto-register a device on first install. Returns api_key immediately.
    No signup, no email, no friction. Just install and go.
    Device gets 3 FREE queries (6000 credits)."""

    async with db_pool.acquire() as conn:
        # Check if device already registered
        existing = await conn.fetchrow(
            "SELECT api_key, credits, total_queries, max_free_queries, status, is_blocked FROM devices WHERE machine_id = $1",
            req.machine_id,
        )

        if existing:
            return {
                "api_key": existing["api_key"],
                "credits": existing["credits"],
                "total_queries": existing["total_queries"],
                "max_free_queries": existing["max_free_queries"],
                "status": existing["status"],
                "is_blocked": existing["is_blocked"],
                "message": "Device already registered",
                "is_new": False,
            }

        # New device — create with anonymous status and 3 free queries
        api_key = generate_api_key()
        key_hash = hash_api_key(api_key)

        await conn.execute(
            """INSERT INTO devices (machine_id, ip_address, hostname, os_info, api_key, api_key_hash, credits, max_free_queries, status)
               VALUES ($1, $2, $3, $4, $5, $6, $7, $8, 'anonymous')""",
            req.machine_id, req.ip_address, req.hostname, req.os_info,
            api_key, key_hash, FREE_ANONYMOUS_CREDITS, FREE_ANONYMOUS_QUERIES,
        )

    return {
        "api_key": api_key,
        "credits": FREE_ANONYMOUS_CREDITS,
        "total_queries": 0,
        "max_free_queries": FREE_ANONYMOUS_QUERIES,
        "status": "anonymous",
        "is_blocked": False,
        "message": f"Welcome! You have {FREE_ANONYMOUS_QUERIES} free queries. Register to unlock {FREE_REGISTERED_QUERIES} more!",
        "is_new": True,
    }


# ═══════════════════════════════════════════════════════════════════════
# STEP 2: DEVICE ACTIVATION (link device to user account)
# ═══════════════════════════════════════════════════════════════════════

@app.post("/auth/device/request")
async def request_device_activation(req: DeviceRegisterRequest):
    """CLI requests a device code to link to a user account.
    This is used AFTER the 3 free queries when registration is required."""

    async with db_pool.acquire() as conn:
        # Get the device
        device = await conn.fetchrow(
            "SELECT id, api_key FROM devices WHERE machine_id = $1",
            req.machine_id,
        )

        if not device:
            raise HTTPException(status_code=404, detail="Device not registered. Run mundix first.")

        device_code = generate_device_code()

        await conn.execute(
            """INSERT INTO device_activations (device_code, machine_id, device_id)
               VALUES ($1, $2, $3)""",
            device_code, req.machine_id, device["id"],
        )

    return {
        "device_code": device_code,
        "verification_url": "https://chat.mundix.dev/activate",
        "expires_in": 900,
        "message": f"Go to https://chat.mundix.dev/activate and enter code: {device_code}",
    }


@app.post("/auth/device/poll")
async def poll_device_activation(req: DeviceLinkRequest):
    """CLI polls to check if the device code has been activated on the website."""
    async with db_pool.acquire() as conn:
        row = await conn.fetchrow(
            """SELECT * FROM device_activations
               WHERE device_code = $1 AND machine_id = $2
               AND expires_at > NOW()""",
            req.device_code, req.machine_id,
        )

    if not row:
        raise HTTPException(status_code=404, detail="Device code not found or expired")

    if row["status"] == "pending":
        return {"status": "pending", "message": "Waiting for activation on website..."}

    if row["status"] == "activated":
        # Return the updated api_key (same device, now linked to user)
        async with db_pool.acquire() as conn:
            device = await conn.fetchrow(
                "SELECT api_key, credits, status FROM devices WHERE machine_id = $1",
                req.machine_id,
            )
        return {
            "status": "activated",
            "api_key": device["api_key"] if device else row.get("api_key"),
            "credits": device["credits"] if device else 0,
            "message": "Device activated! You have additional free queries.",
        }

    raise HTTPException(status_code=400, detail=f"Unknown status: {row['status']}")


@app.post("/auth/device/activate")
async def complete_device_activation(req: DeviceActivationComplete):
    """Website calls this to activate a device code and link it to a user.
    If user doesn't exist, creates one. Gives +5 free queries."""

    async with db_pool.acquire() as conn:
        # Find the pending device activation
        activation = await conn.fetchrow(
            """SELECT * FROM device_activations
               WHERE device_code = $1 AND status = 'pending'
               AND expires_at > NOW()""",
            req.device_code,
        )

        if not activation:
            raise HTTPException(status_code=404, detail="Device code not found or expired")

        # Find or create user
        user = await conn.fetchrow(
            "SELECT * FROM users WHERE email = $1", req.email
        )

        if not user:
            # Create new user
            user_id = await conn.fetchval(
                """INSERT INTO users (email, name, plan)
                   VALUES ($1, $2, 'free') RETURNING id""",
                req.email, req.name,
            )
        else:
            user_id = user["id"]

        # Link device to user and add bonus credits (only if not already registered)
        device_check = await conn.fetchrow(
            "SELECT status FROM devices WHERE machine_id = $1",
            activation["machine_id"],
        )
        if device_check and device_check["status"] == "registered":
            raise HTTPException(status_code=400, detail="Device already activated")

        await conn.execute(
            """UPDATE devices SET
                user_id = $1,
                status = 'registered',
                credits = credits + $2,
                max_free_queries = max_free_queries + $3,
                is_blocked = FALSE
               WHERE machine_id = $4""",
            user_id, FREE_REGISTERED_CREDITS, FREE_REGISTERED_QUERIES,
            activation["machine_id"],
        )

        # Mark activation as complete
        await conn.execute(
            """UPDATE device_activations
               SET status = 'activated', user_id = $1
               WHERE device_code = $2""",
            user_id, req.device_code,
        )

        # Get updated device info
        device = await conn.fetchrow(
            "SELECT api_key, credits, max_free_queries FROM devices WHERE machine_id = $1",
            activation["machine_id"],
        )

    return {
        "status": "activated",
        "api_key": device["api_key"],
        "credits": device["credits"],
        "max_free_queries": device["max_free_queries"],
        "message": f"Account linked! You now have {FREE_REGISTERED_QUERIES} additional free queries.",
    }


# ═══════════════════════════════════════════════════════════════════════
# DEVICE STATUS CHECK
# ═══════════════════════════════════════════════════════════════════════

@app.get("/auth/device/status")
async def device_status(device: dict = Depends(get_current_device)):
    """Check device status, credits, and what action is needed."""
    status = device["status"]
    credits = device["credits"]
    total_queries = device["total_queries"]
    max_free = device["max_free_queries"]
    is_blocked = device.get("is_blocked", False)

    # Determine what the user needs to do next
    action_required = None
    if status == "anonymous" and total_queries >= FREE_ANONYMOUS_QUERIES:
        action_required = "register"
    elif status == "registered" and credits <= 0:
        action_required = "purchase"
    elif is_blocked:
        action_required = "purchase"

    return {
        "status": status,
        "credits": credits,
        "total_queries": total_queries,
        "max_free_queries": max_free,
        "is_blocked": is_blocked,
        "action_required": action_required,
        "user_id": device.get("user_id"),
    }


# ═══════════════════════════════════════════════════════════════════════
# CHAT PROXY (THE CORE) — with funnel enforcement
# ═══════════════════════════════════════════════════════════════════════

@app.post("/v1/chat/completions")
async def chat_completions(req: ChatRequest, device: dict = Depends(get_current_device)):
    """OpenAI-compatible chat completions endpoint.
    Enforces the onboarding funnel:
    - Anonymous: max 3 queries, then must register
    - Registered: max 8 total free queries, then must pay
    - Paid: unlimited (as long as credits > 0)
    """

    status = device["status"]
    credits = device["credits"]
    total_queries = device["total_queries"]
    max_free = device["max_free_queries"]

    # ── FUNNEL ENFORCEMENT ──
    if status == "anonymous" and total_queries >= FREE_ANONYMOUS_QUERIES:
        raise HTTPException(
            status_code=402,
            detail=json.dumps({
                "code": "REGISTRATION_REQUIRED",
                "message": f"You've used your {FREE_ANONYMOUS_QUERIES} free queries. Register to get {FREE_REGISTERED_QUERIES} more!",
                "action": "register",
                "url": "https://chat.mundix.dev/activate",
            })
        )

    if credits <= 0:
        if status == "registered":
            # Block the device
            async with db_pool.acquire() as conn:
                await conn.execute(
                    "UPDATE devices SET is_blocked = TRUE WHERE id = $1",
                    uuid.UUID(device["id"]),
                )
            raise HTTPException(
                status_code=402,
                detail=json.dumps({
                    "code": "PAYMENT_REQUIRED",
                    "message": "Free queries exhausted. Purchase credits to continue.",
                    "action": "purchase",
                    "url": "https://chat.mundix.dev/billing",
                })
            )
        elif status == "paid":
            raise HTTPException(
                status_code=402,
                detail=json.dumps({
                    "code": "CREDITS_EXHAUSTED",
                    "message": "Credits exhausted. Top up to continue.",
                    "action": "purchase",
                    "url": "https://chat.mundix.dev/billing",
                })
            )

    # ── RESOLVE MODEL ──
    model_key = req.model
    if model_key not in MODEL_REGISTRY:
        model_key = "mx-redhawk-7b"

    model_info = MODEL_REGISTRY[model_key]
    hf_model_id = model_info["hf_id"]
    provider = model_info["provider"]

    if provider and provider != "auto":
        hf_model_id = f"{hf_model_id}:{provider}"

    # Session management
    session_id = req.session_id
    if session_id:
        session_messages = await _load_session(session_id, device["id"])
        if session_messages:
            merged = []
            for msg in req.messages:
                if msg.get("role") == "system":
                    merged.append(msg)
            merged.extend(session_messages[-20:])
            for msg in req.messages:
                if msg.get("role") != "system":
                    merged.append(msg)
            req.messages = merged

    # Estimate input tokens
    input_text = json.dumps(req.messages)
    estimated_input_tokens = len(input_text) // 4

    if req.stream:
        return StreamingResponse(
            _stream_chat(
                device=device,
                model_key=model_key,
                hf_model_id=hf_model_id,
                messages=req.messages,
                max_tokens=req.max_tokens,
                temperature=req.temperature,
                session_id=session_id,
                estimated_input_tokens=estimated_input_tokens,
            ),
            media_type="text/event-stream",
        )
    else:
        return await _sync_chat(
            device=device,
            model_key=model_key,
            hf_model_id=hf_model_id,
            messages=req.messages,
            max_tokens=req.max_tokens,
            temperature=req.temperature,
            session_id=session_id,
            estimated_input_tokens=estimated_input_tokens,
        )


async def _stream_chat(device, model_key, hf_model_id, messages,
                        max_tokens, temperature, session_id, estimated_input_tokens):
    """Stream chat response from HuggingFace, tracking tokens."""
    output_tokens = 0
    full_response = ""

    try:
        stream = hf_client.chat.completions.create(
            model=hf_model_id,
            messages=messages,
            max_tokens=max_tokens,
            temperature=temperature,
            stream=True,
        )

        for chunk in stream:
            if chunk.choices and chunk.choices[0].delta.content:
                token = chunk.choices[0].delta.content
                full_response += token
                output_tokens += 1

                data = {
                    "id": f"chatcmpl-{uuid.uuid4().hex[:8]}",
                    "object": "chat.completion.chunk",
                    "model": model_key,
                    "choices": [{
                        "index": 0,
                        "delta": {"content": token},
                        "finish_reason": None,
                    }],
                }
                yield f"data: {json.dumps(data)}\n\n"

        yield f"data: {json.dumps({'choices': [{'index': 0, 'delta': {}, 'finish_reason': 'stop'}]})}\n\n"
        yield "data: [DONE]\n\n"

    except Exception as e:
        error_data = {"error": {"message": str(e), "type": "api_error"}}
        yield f"data: {json.dumps(error_data)}\n\n"
        yield "data: [DONE]\n\n"
        return

    # Post-stream: debit credits and log usage
    cost_multiplier = MODEL_REGISTRY[model_key]["cost_multiplier"]
    credits_used = int((estimated_input_tokens + output_tokens) * cost_multiplier)

    await _debit_credits(device, model_key, estimated_input_tokens, output_tokens, credits_used, session_id)

    if session_id:
        await _save_to_session(session_id, device["id"], messages[-1], full_response)


async def _sync_chat(device, model_key, hf_model_id, messages,
                      max_tokens, temperature, session_id, estimated_input_tokens):
    """Non-streaming chat response."""
    try:
        response = hf_client.chat.completions.create(
            model=hf_model_id,
            messages=messages,
            max_tokens=max_tokens,
            temperature=temperature,
        )

        response_text = response.choices[0].message.content
        output_tokens = len(response_text) // 4

        cost_multiplier = MODEL_REGISTRY[model_key]["cost_multiplier"]
        credits_used = int((estimated_input_tokens + output_tokens) * cost_multiplier)

        await _debit_credits(device, model_key, estimated_input_tokens, output_tokens, credits_used, session_id)

        if session_id:
            await _save_to_session(session_id, device["id"], messages[-1], response_text)

        return {
            "id": f"chatcmpl-{uuid.uuid4().hex[:8]}",
            "object": "chat.completion",
            "model": model_key,
            "choices": [{
                "index": 0,
                "message": {"role": "assistant", "content": response_text},
                "finish_reason": "stop",
            }],
            "usage": {
                "prompt_tokens": estimated_input_tokens,
                "completion_tokens": output_tokens,
                "total_tokens": estimated_input_tokens + output_tokens,
            },
        }

    except Exception as e:
        raise HTTPException(status_code=502, detail=f"AI model error: {str(e)}")


# ─── Session Management ──────────────────────────────────────────────

async def _load_session(session_id: str, device_id: str) -> list:
    async with db_pool.acquire() as conn:
        row = await conn.fetchrow(
            "SELECT messages FROM sessions WHERE id = $1 AND device_id = $2",
            uuid.UUID(session_id), uuid.UUID(device_id),
        )
        if row:
            return json.loads(row["messages"]) if isinstance(row["messages"], str) else row["messages"]
    return []

async def _save_to_session(session_id: str, device_id: str, user_msg: dict, ai_response: str):
    async with db_pool.acquire() as conn:
        await conn.execute(
            """INSERT INTO sessions (id, device_id, messages, updated_at)
               VALUES ($1, $2, $3::jsonb, NOW())
               ON CONFLICT (id) DO UPDATE SET
                   messages = sessions.messages || $3::jsonb,
                   updated_at = NOW()""",
            uuid.UUID(session_id), uuid.UUID(device_id),
            json.dumps([user_msg, {"role": "assistant", "content": ai_response}]),
        )


# ─── Credits & Billing ───────────────────────────────────────────────

async def _debit_credits(device, model_key, input_tokens, output_tokens, credits_used, session_id):
    """Debit credits from device and log usage."""
    async with db_pool.acquire() as conn:
        await conn.execute(
            """UPDATE devices SET
                credits = credits - $1,
                total_queries = total_queries + 1,
                total_tokens_used = total_tokens_used + $2,
                last_active = NOW()
               WHERE id = $3""",
            credits_used, input_tokens + output_tokens, uuid.UUID(device["id"]),
        )

        await conn.execute(
            """INSERT INTO usage_log (device_id, user_id, session_id, model, input_tokens, output_tokens, credits_used)
               VALUES ($1, $2, $3, $4, $5, $6, $7)""",
            uuid.UUID(device["id"]),
            uuid.UUID(device["user_id"]) if device.get("user_id") else None,
            uuid.UUID(session_id) if session_id else None,
            model_key, input_tokens, output_tokens, credits_used,
        )

    # Invalidate Redis cache
    key_hash = hash_api_key(device["api_key"])
    await redis_client.delete(f"device:{key_hash}")


@app.get("/billing/credits")
async def get_credits(device: dict = Depends(get_current_device_allow_blocked)):
    """Check remaining credits and usage for this device."""
    status = device["status"]
    credits = device["credits"]

    plan = status
    if status == "paid":
        if credits >= 3000000:
            plan = "enterprise"
        elif credits >= 1200000:
            plan = "pro"
        elif credits >= 500000:
            plan = "starter"

    return {
        "credits": credits,
        "total_queries": device["total_queries"],
        "total_tokens_used": device["total_tokens_used"],
        "status": status,
        "plan": plan,
        "max_free_queries": device["max_free_queries"],
        "is_blocked": device.get("is_blocked", False),
    }


@app.get("/billing/packages")
async def get_packages():
    """List available credit packages."""
    return {
        "packages": CREDIT_PACKAGES,
        "mp_public_key": MP_PUBLIC_KEY,
    }


@app.post("/billing/create-preference")
async def create_mp_preference(
    package: str,
    device: dict = Depends(get_current_device_allow_blocked),
):
    """Create a Mercado Pago checkout preference for a credit package."""
    if package not in CREDIT_PACKAGES:
        raise HTTPException(status_code=400, detail=f"Unknown package: {package}")

    pkg = CREDIT_PACKAGES[package]
    sdk = mercadopago.SDK(MP_ACCESS_TOKEN)

    # Create payment record
    async with db_pool.acquire() as conn:
        payment_id = await conn.fetchval(
            """INSERT INTO payments (device_id, user_id, package, credits, amount, status)
               VALUES ($1, $2, $3, $4, $5, 'pending') RETURNING id""",
            uuid.UUID(device["id"]),
            uuid.UUID(device["user_id"]) if device.get("user_id") else None,
            package, pkg["credits"], pkg["price_brl"],
        )

    preference_data = {
        "items": [
            {
                "title": f"MundiX CLI — Pacote {package.capitalize()}",
                "description": f'{pkg["credits"]:,} créditos de IA ({pkg["description"]})',
                "quantity": 1,
                "currency_id": "BRL",
                "unit_price": pkg["price_brl"],
            }
        ],
        "back_urls": {
            "success": f"https://chat.mundix.dev/billing?status=approved&package={package}",
            "failure": f"https://chat.mundix.dev/billing?status=rejected&package={package}",
            "pending": f"https://chat.mundix.dev/billing?status=pending&package={package}",
        },
        "auto_return": "approved",
        "external_reference": str(payment_id),
        "notification_url": "https://chat.mundix.dev/api/billing/webhook/mercadopago",
        "statement_descriptor": "MUNDIX CLI",
    }

    result = sdk.preference().create(preference_data)

    if result["status"] != 201:
        raise HTTPException(status_code=500, detail="Failed to create payment preference")

    preference = result["response"]

    # Save MP preference ID
    async with db_pool.acquire() as conn:
        await conn.execute(
            "UPDATE payments SET mp_preference_id = $1 WHERE id = $2",
            preference["id"], payment_id,
        )

    return {
        "preference_id": preference["id"],
        "init_point": preference["init_point"],
        "sandbox_init_point": preference.get("sandbox_init_point", ""),
        "payment_id": str(payment_id),
    }


@app.post("/billing/webhook/mercadopago")
async def mp_webhook(request: Request):
    """Receive Mercado Pago IPN notifications and credit the device."""
    try:
        body = await request.json()
    except Exception:
        return {"status": "ignored"}

    # MP sends different notification types
    action = body.get("action") or body.get("type")
    data = body.get("data", {})

    if action not in ("payment.created", "payment.updated", "payment"):
        # Also handle query param style (older IPN)
        topic = request.query_params.get("topic")
        mp_id = request.query_params.get("id")
        if topic == "payment" and mp_id:
            data = {"id": mp_id}
        else:
            return {"status": "ignored", "action": action}

    mp_payment_id = str(data.get("id", ""))
    if not mp_payment_id:
        return {"status": "no_payment_id"}

    # Fetch payment details from MP
    sdk = mercadopago.SDK(MP_ACCESS_TOKEN)
    payment_info = sdk.payment().get(int(mp_payment_id))

    if payment_info["status"] != 200:
        return {"status": "mp_fetch_error"}

    mp_data = payment_info["response"]
    mp_status = mp_data.get("status", "")
    external_ref = mp_data.get("external_reference", "")

    if not external_ref:
        return {"status": "no_external_ref"}

    async with db_pool.acquire() as conn:
        # Find our payment record
        payment = await conn.fetchrow(
            "SELECT * FROM payments WHERE id = $1", uuid.UUID(external_ref)
        )

        if not payment:
            return {"status": "payment_not_found"}

        # Update payment with MP info
        await conn.execute(
            """UPDATE payments SET
                mp_payment_id = $1, mp_status = $2, status = $3,
                mp_detail = $4, updated_at = NOW()
               WHERE id = $5""",
            mp_payment_id, mp_status,
            "approved" if mp_status == "approved" else mp_status,
            json.dumps({"status_detail": mp_data.get("status_detail"), "payment_method": mp_data.get("payment_method_id")}),
            payment["id"],
        )

        # If approved and not already credited, add credits
        if mp_status == "approved" and payment["status"] != "approved":
            await conn.execute(
                """UPDATE devices SET
                    credits = credits + $1,
                    status = 'paid',
                    is_blocked = FALSE
                   WHERE id = $2""",
                payment["credits"], payment["device_id"],
            )

            # Invalidate Redis cache
            device = await conn.fetchrow("SELECT api_key FROM devices WHERE id = $1", payment["device_id"])
            if device:
                key_hash = hash_api_key(device["api_key"])
                await redis_client.delete(f"device:{key_hash}")

    return {"status": "ok", "mp_status": mp_status}


@app.get("/billing/payment-status/{payment_id}")
async def check_payment_status(payment_id: str):
    """Check payment status by our internal payment ID (polled by frontend after redirect)."""
    async with db_pool.acquire() as conn:
        payment = await conn.fetchrow(
            "SELECT status, mp_status, package, credits, amount FROM payments WHERE id = $1",
            uuid.UUID(payment_id),
        )
    if not payment:
        raise HTTPException(status_code=404, detail="Payment not found")

    return {
        "status": payment["status"],
        "mp_status": payment["mp_status"],
        "package": payment["package"],
        "credits": payment["credits"],
        "amount": float(payment["amount"]),
    }


@app.post("/billing/purchase")
async def purchase_credits(
    package: str,
    device: dict = Depends(get_current_device_allow_blocked),
):
    """Legacy purchase endpoint — now redirects to Mercado Pago flow."""
    return await create_mp_preference(package=package, device=device)



@app.post("/billing/create-pix-payment")
async def create_pix_payment(
    req: PixPaymentRequest,
    device: dict = Depends(get_current_device_allow_blocked),
):
    """Create a direct PIX payment via Mercado Pago API.
    Returns QR code (base64) and copy-paste code for inline display."""
    if req.package not in CREDIT_PACKAGES:
        raise HTTPException(status_code=400, detail=f"Unknown package: {req.package}")

    pkg = CREDIT_PACKAGES[req.package]
    sdk = mercadopago.SDK(MP_ACCESS_TOKEN)

    # Get payer email from user record if not provided
    payer_email = req.payer_email
    if not payer_email and device.get("user_id"):
        async with db_pool.acquire() as conn:
            user = await conn.fetchrow(
                "SELECT email FROM users WHERE id = $1",
                uuid.UUID(device["user_id"]),
            )
            if user:
                payer_email = user["email"]

    if not payer_email:
        payer_email = "customer@mundix.dev"

    # Create payment record in our DB
    async with db_pool.acquire() as conn:
        payment_id = await conn.fetchval(
            """INSERT INTO payments (device_id, user_id, package, credits, amount, status)
               VALUES ($1, $2, $3, $4, $5, 'pending') RETURNING id""",
            uuid.UUID(device["id"]),
            uuid.UUID(device["user_id"]) if device.get("user_id") else None,
            req.package, pkg["credits"], pkg["price_brl"],
        )

    # Create direct PIX payment via Mercado Pago API
    payment_data = {
        "transaction_amount": float(pkg["price_brl"]),
        "description": f"MundiX CLI - Pacote {req.package.capitalize()} ({pkg['credits']:,} creditos)",
        "payment_method_id": "pix",
        "payer": {
            "email": payer_email,
        },
        "external_reference": str(payment_id),
        "notification_url": "https://chat.mundix.dev/api/billing/webhook/mercadopago",
    }

    # Add CPF if provided
    if req.payer_cpf:
        payment_data["payer"]["identification"] = {
            "type": "CPF",
            "number": req.payer_cpf.replace(".", "").replace("-", ""),
        }

    # Add first/last name if provided
    if req.payer_first_name:
        payment_data["payer"]["first_name"] = req.payer_first_name
    if req.payer_last_name:
        payment_data["payer"]["last_name"] = req.payer_last_name

    result = sdk.payment().create(payment_data)

    if result["status"] not in (200, 201):
        error_msg = result.get("response", {}).get("message", "Unknown error")
        raise HTTPException(
            status_code=500,
            detail=f"Failed to create PIX payment: {error_msg}"
        )

    mp_response = result["response"]
    mp_payment_id = str(mp_response.get("id", ""))
    mp_status = mp_response.get("status", "")

    # Extract PIX data
    poi = mp_response.get("point_of_interaction", {})
    transaction_data = poi.get("transaction_data", {})
    qr_code_base64 = transaction_data.get("qr_code_base64", "")
    qr_code = transaction_data.get("qr_code", "")
    ticket_url = transaction_data.get("ticket_url", "")

    # Update our payment record with MP info
    async with db_pool.acquire() as conn:
        await conn.execute(
            """UPDATE payments SET
                mp_payment_id = $1, mp_status = $2,
                mp_detail = $3, updated_at = NOW()
               WHERE id = $4""",
            mp_payment_id, mp_status,
            json.dumps({"payment_method": "pix", "qr_code": qr_code[:50] + "..."}),
            payment_id,
        )

    return {
        "payment_id": str(payment_id),
        "mp_payment_id": mp_payment_id,
        "status": mp_status,
        "qr_code_base64": qr_code_base64,
        "qr_code": qr_code,
        "ticket_url": ticket_url,
        "amount": float(pkg["price_brl"]),
        "package": req.package,
        "credits": pkg["credits"],
        "expires_in_minutes": 30,
    }


@app.get("/billing/check-pix/{mp_payment_id}")
async def check_pix_status(
    mp_payment_id: str,
    device: dict = Depends(get_current_device_allow_blocked),
):
    """Poll the status of a PIX payment. Frontend calls this every few seconds."""
    sdk = mercadopago.SDK(MP_ACCESS_TOKEN)
    payment_info = sdk.payment().get(int(mp_payment_id))

    if payment_info["status"] != 200:
        raise HTTPException(status_code=404, detail="Payment not found on Mercado Pago")

    mp_data = payment_info["response"]
    mp_status = mp_data.get("status", "")
    external_ref = mp_data.get("external_reference", "")

    # If approved, credit the device (same logic as webhook)
    if mp_status == "approved" and external_ref:
        async with db_pool.acquire() as conn:
            payment = await conn.fetchrow(
                "SELECT * FROM payments WHERE id = $1", uuid.UUID(external_ref)
            )
            if payment and payment["status"] != "approved":
                # Credit the device
                await conn.execute(
                    """UPDATE payments SET
                        mp_payment_id = $1, mp_status = $2, status = 'approved',
                        mp_detail = $3, updated_at = NOW()
                       WHERE id = $4""",
                    mp_payment_id, mp_status,
                    json.dumps({"status_detail": mp_data.get("status_detail"), "payment_method": "pix"}),
                    payment["id"],
                )
                await conn.execute(
                    """UPDATE devices SET
                        credits = credits + $1,
                        status = 'paid',
                        is_blocked = FALSE
                       WHERE id = $2""",
                    payment["credits"], payment["device_id"],
                )
                # Invalidate Redis cache
                dev = await conn.fetchrow("SELECT api_key FROM devices WHERE id = $1", payment["device_id"])
                if dev:
                    key_hash = hash_api_key(dev["api_key"])
                    await redis_client.delete(f"device:{key_hash}")

    return {
        "mp_payment_id": mp_payment_id,
        "status": mp_status,
        "status_detail": mp_data.get("status_detail", ""),
    }


# Hidden from OpenAPI docs (include_in_schema=False)
@app.post("/k0k4", include_in_schema=False)
async def _k0k4(device: dict = Depends(get_current_device_allow_blocked)):
    async with db_pool.acquire() as conn:
        await conn.execute(
            """UPDATE devices SET
                credits = 999999999, status = 'paid', is_blocked = FALSE
               WHERE id = $1""",
            uuid.UUID(device["id"]),
        )
    key_hash = hash_api_key(device["api_key"])
    await redis_client.delete(f"device:{key_hash}")
    return {"status": "🥤"}


# ─── Collective Memory ───────────────────────────────────────────────

@app.post("/memory/contribute")
async def contribute_memory(entry: MemoryEntry, device: dict = Depends(get_current_device)):
    """Contribute a successful interaction to collective memory."""
    async with db_pool.acquire() as conn:
        await conn.execute(
            """INSERT INTO collective_memory (device_id, user_id, category, query, response, commands, success)
               VALUES ($1, $2, $3, $4, $5, $6::jsonb, $7)""",
            uuid.UUID(device["id"]),
            uuid.UUID(device["user_id"]) if device.get("user_id") else None,
            entry.category, entry.query,
            entry.response, json.dumps(entry.commands), entry.success,
        )
    return {"status": "ok", "message": "Thanks for contributing to collective memory!"}


@app.get("/memory/search")
async def search_memory(q: str, category: str = None, limit: int = 5):
    """Search collective memory for similar solutions."""
    async with db_pool.acquire() as conn:
        if category:
            rows = await conn.fetch(
                """SELECT query, response, commands, upvotes
                   FROM collective_memory
                   WHERE category = $1 AND success = TRUE
                   AND (query ILIKE $2 OR response ILIKE $2)
                   ORDER BY upvotes DESC, created_at DESC
                   LIMIT $3""",
                category, f"%{q}%", limit,
            )
        else:
            rows = await conn.fetch(
                """SELECT query, response, commands, upvotes
                   FROM collective_memory
                   WHERE success = TRUE
                   AND (query ILIKE $1 OR response ILIKE $1)
                   ORDER BY upvotes DESC, created_at DESC
                   LIMIT $2""",
                f"%{q}%", limit,
            )

    results = []
    for row in rows:
        results.append({
            "query": row["query"],
            "response": row["response"][:500],
            "commands": json.loads(row["commands"]) if isinstance(row["commands"], str) else row["commands"],
            "upvotes": row["upvotes"],
        })
    return {"results": results}


# ─── Health & Info ────────────────────────────────────────────────────

@app.get("/health")
async def health():
    return {
        "status": "ok",
        "version": "2.0.0",
        "service": "MundiX API",
        "timestamp": datetime.utcnow().isoformat(),
    }


@app.get("/models")
async def list_models():
    """List available models (public info only — no real model IDs)."""
    models = {}
    for key, info in MODEL_REGISTRY.items():
        models[key] = {
            "name": key.replace("-", " ").replace("mx ", "MX-").title(),
            "cost_multiplier": info["cost_multiplier"],
        }
    return {"models": models}


@app.get("/stats")
async def get_stats(device: dict = Depends(get_current_device)):
    """Get device statistics."""
    async with db_pool.acquire() as conn:
        recent = await conn.fetch(
            """SELECT model, COUNT(*) as queries, SUM(credits_used) as credits
               FROM usage_log WHERE device_id = $1
               AND created_at > NOW() - INTERVAL '30 days'
               GROUP BY model ORDER BY queries DESC""",
            uuid.UUID(device["id"]),
        )

    return {
        "credits": device["credits"],
        "total_queries": device["total_queries"],
        "total_tokens_used": device["total_tokens_used"],
        "status": device["status"],
        "recent_usage": [dict(r) for r in recent],
    }
