import os
import re
import json
import time
import uuid
import asyncio
import logging
import tempfile
import threading
from pathlib import Path
from datetime import datetime
from concurrent.futures import ThreadPoolExecutor, TimeoutError as FuturesTimeoutError
from logging.handlers import RotatingFileHandler

from curl_cffi import requests as cffi_requests
from perplexity import Client
from perplexity.config import DEFAULT_HEADERS, EMAILNATOR_BASE_URL, EMAILNATOR_HEADERS
from perplexity.exceptions import PerplexityError, AccountCreationError
from perplexity.utils import retry_with_backoff, sanitize_query
from perplexity.logger import get_logger
from config_models import MODEL_MAPPINGS, LABS_MODELS, MODE_MAP, VALID_MODEL_NAMES

from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import StreamingResponse, JSONResponse, HTMLResponse
from fastapi.middleware.cors import CORSMiddleware

logger = get_logger("app")

_data_candidate = Path(os.environ.get("DATA_DIR", "/data"))
DATA_DIR = _data_candidate if _data_candidate.is_dir() and os.access(_data_candidate, os.W_OK) else Path(tempfile.gettempdir()) / "ws"
DATA_DIR.mkdir(parents=True, exist_ok=True)

COOKIES_FILE = DATA_DIR / "ws_cookies.json"
EMAILNATOR_COOKIES_FILE = DATA_DIR / "emailnator_cookies.json"
HISTORY_FILE = DATA_DIR / "ws_history.json"
LOG_FILE = DATA_DIR / "perplexity.log"
THREADS_FILE = DATA_DIR / "ws_threads.json"
START_TIME = time.time()

LOCK = threading.Lock()
ACCOUNT_LOCK = threading.Lock()
EXECUTOR = ThreadPoolExecutor(max_workers=2)
ACCOUNT_CREATE_TIMEOUT = 120

file_handler = RotatingFileHandler(str(LOG_FILE), maxBytes=5*1024*1024, backupCount=3)
file_handler.setFormatter(logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s"))
logging.getLogger("perplexity").addHandler(file_handler)
logging.getLogger("app").addHandler(file_handler)

DEFAULT_MODE = "pro"


def _content_to_str(content):
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        return " ".join(b.get("text", "") for b in content if isinstance(b, dict) and b.get("type") in ("text", "input_text"))
    return str(content)


MAX_HISTORY = 50

def _load_history():
    if HISTORY_FILE.exists():
        try:
            with open(HISTORY_FILE, "r") as f:
                return json.load(f)
        except Exception:
            pass
    return []

def _add_history(entry):
    with LOCK:
        h = _load_history()
        h.append(entry)
        h = h[-MAX_HISTORY:]
        with open(HISTORY_FILE, "w") as f:
            json.dump(h, f)

# ─── Thread Store ────────────────────────────────────────────────────────────

def _load_threads():
    if THREADS_FILE.exists():
        try:
            with open(THREADS_FILE, "r") as f:
                return json.load(f)
        except Exception:
            pass
    return {}

def _save_threads(threads):
    with LOCK:
        with open(THREADS_FILE, "w") as f:
            json.dump(threads, f)

def _get_thread(thread_id):
    return _load_threads().get(thread_id)

def _update_thread(thread_id, data):
    threads = _load_threads()
    if thread_id not in threads:
        threads[thread_id] = {"id": thread_id, "created_at": datetime.utcnow().isoformat(), "messages": [], "backend_uuid": None, "frontend_context_uuid": None, "mode": DEFAULT_MODE}
    threads[thread_id].update(data)
    _save_threads(threads)
    return threads[thread_id]


# ─── Cookie / Account Management ─────────────────────────────────────────────

def load_cookies():
    if COOKIES_FILE.exists():
        try:
            with open(COOKIES_FILE, "r") as f:
                data = json.load(f)
            if data.get("cookies") and time.time() - data.get("saved_at", 0) < 86400:
                return data
        except Exception:
            pass
    return None

def save_cookies(cookies, copilot_remaining=5, file_upload_remaining=10):
    with LOCK:
        with open(COOKIES_FILE, "w") as f:
            json.dump({"cookies": cookies, "saved_at": time.time(), "copilot_remaining": copilot_remaining, "file_upload_remaining": file_upload_remaining}, f)

def _load_emailnator_cookies():
    if EMAILNATOR_COOKIES_FILE.exists():
        try:
            with open(EMAILNATOR_COOKIES_FILE, "r") as f:
                return json.load(f)
        except Exception:
            pass
    return {}

def _save_emailnator_cookies(cookies):
    with LOCK:
        with open(EMAILNATOR_COOKIES_FILE, "w") as f:
            json.dump(cookies, f)

def _init_emailnator_session():
    cookies = _load_emailnator_cookies()
    session = cffi_requests.Session(headers=EMAILNATOR_HEADERS.copy(), impersonate="chrome")
    if cookies:
        for k, v in cookies.items():
            session.cookies.set(k, v)
    try:
        session.get(EMAILNATOR_BASE_URL, timeout=10)
        _save_emailnator_cookies(session.cookies.get_dict())
    except Exception:
        pass
    return session

def _get_csrf_token(session):
    try:
        resp = session.get("https://www.perplexity.ai/api/auth/csrf", timeout=15)
        if resp.ok:
            return resp.json().get("csrfToken", "")
    except Exception:
        pass
    for cookie in session.cookies:
        if cookie.name == "__Host-next-auth.csrf-token":
            val = cookie.value.split("%7C")[0] if "%7C" in cookie.value else cookie.value
            return val
    return ""

def _extract_signin_link(email_body):
    pattern = r'https://www\.perplexity\.ai/api/auth/callback/email\?[^"\s<>]+'
    match = re.search(pattern, email_body)
    if match:
        link = match.group(0)
        if link.endswith("="):
            link = link[:-1]
        return link
    return None

def _create_account_emailnator(session, csrf):
    from perplexity.emailnator import Emailnator
    emailnator_cli = Emailnator()
    emailnator_cli.get_email()
    email = emailnator_cli.email
    if not email:
        return None
    signin_resp = session.post("https://www.perplexity.ai/api/auth/signin/email", json={"email": email, "csrfToken": csrf}, headers={"content-type": "application/json"}, timeout=15)
    if not signin_resp.ok:
        return None
    try:
        msg = emailnator_cli.get(func=lambda x: x["subject"] == "Sign in to Perplexity")
        link = _extract_signin_link(emailnator_cli.open(msg["messageID"]))
        return link
    except Exception:
        return None

def _create_account_guerrilla(session, csrf):
    try:
        gu_session = cffi_requests.Session(impersonate="chrome", timeout=15)
        gu_resp = gu_session.get("https://api.guerrillamail.com/ajax.php?f=get_email_address")
        if not gu_resp.ok:
            return None
        email = gu_resp.json().get("email_addr", "")
        if not email:
            return None
        signin_resp = session.post("https://www.perplexity.ai/api/auth/signin/email", json={"email": email, "csrfToken": csrf}, headers={"content-type": "application/json"}, timeout=15)
        if not signin_resp.ok:
            return None
        sid_token = gu_resp.json().get("sid_token", "")
        for attempt in range(8):
            time.sleep(3)
            check = gu_session.get(f"https://api.guerrillamail.com/ajax.php?f=check_email&sid_token={sid_token}&seq=0")
            if not check.ok:
                continue
            emails = check.json().get("list", [])
            for em in emails:
                if "Perplexity" in em.get("mail_subject", ""):
                    mail_body = gu_session.get(f"https://api.guerrillamail.com/ajax.php?f=fetch_email&sid_token={sid_token}&email_id={em['mail_id']}").json().get("mail_body", "")
                    link = _extract_signin_link(mail_body)
                    if link:
                        return link
        return None
    except Exception as e:
        logger.warning(f"Guerrilla Mail failed: {e}")
        return None

@retry_with_backoff(max_attempts=2, exceptions=(AccountCreationError,))
def auto_create_account():
    with ACCOUNT_LOCK:
        session = cffi_requests.Session(headers=DEFAULT_HEADERS.copy(), impersonate="chrome")
        session.get("https://www.perplexity.ai/api/auth/session", timeout=15)
        csrf = _get_csrf_token(session)
        if not csrf:
            raise AccountCreationError("No CSRF token available")
        signin_link = _create_account_guerrilla(session, csrf)
        if not signin_link:
            logger.info("Guerrilla Mail failed, trying Emailnator fallback...")
            session.get("https://www.perplexity.ai/api/auth/session", timeout=15)
            csrf = _get_csrf_token(session)
            if not csrf:
                raise AccountCreationError("No CSRF token available (fallback)")
            signin_link = _create_account_emailnator(session, csrf)
        if not signin_link:
            raise AccountCreationError("Both Guerrilla Mail and Emailnator failed")
        session.get(signin_link, timeout=15)
        final_cookies = {k: v for k, v in session.cookies.get_dict().items()}
        return final_cookies

def _get_client(mode=None):
    stored = load_cookies()
    if stored and stored.get("cookies"):
        client = Client(stored["cookies"])
        client.copilot = stored.get("copilot_remaining", 5)
        client.file_upload = stored.get("file_upload_remaining", 10)
        if client.copilot > 0:
            return client
        logger.info("Copilot credits exhausted, creating new account...")
    logger.info("Creating new account...")
    future = EXECUTOR.submit(auto_create_account)
    try:
        cookies = future.result(timeout=ACCOUNT_CREATE_TIMEOUT)
    except FuturesTimeoutError:
        raise AccountCreationError(f"Account creation timed out after {ACCOUNT_CREATE_TIMEOUT}s")
    client = Client(cookies)
    save_cookies(cookies, client.copilot, client.file_upload)
    return client


# ─── Search Functions ────────────────────────────────────────────────────────

def search_sync(query, mode=None, model=None, sources="web", language="en-US", thread_id=None):
    query = sanitize_query(query)
    sources_list = [s.strip() for s in sources.split(",")] if sources else ["web"]
    mode = mode or DEFAULT_MODE
    model_val = model if model and model in VALID_MODEL_NAMES else None
    thread = _get_thread(thread_id) if thread_id else None
    follow_up = None
    t0 = time.time()
    error = None

    if thread and thread.get("backend_uuid"):
        follow_up = {"backend_uuid": thread["backend_uuid"], "frontend_context_uuid": thread.get("frontend_context_uuid", str(uuid.uuid4())), "attachments": []}
        logger.info(f"Using follow_up for thread {thread_id}")

    try:
        client = _get_client(mode)
        resp = client.search(query=query, mode=mode, model=model_val, sources=sources_list, language=language, follow_up=follow_up)
        save_cookies(client.session.cookies.get_dict() if hasattr(client.session, "cookies") else {}, client.copilot, client.file_upload)
        answer = resp.get("answer", "No answer returned")
    except Exception as e:
        error = str(e)
        answer = None

    elapsed = round(time.time() - t0, 2)
    _add_history({
        "ts": datetime.utcnow().isoformat(timespec="seconds"),
        "query": query[:200],
        "mode": mode,
        "model": model_val,
        "success": error is None,
        "elapsed": elapsed,
        "thread_id": thread_id
    })

    if error:
        raise Exception(error)

    if thread_id:
        backend_uuid = resp.get("backend_uuid") or resp.get("uuid")
        frontend_context_uuid = resp.get("frontend_context_uuid") or str(uuid.uuid4())
        _update_thread(thread_id, {"mode": mode, "backend_uuid": backend_uuid or (thread.get("backend_uuid") if thread else None), "frontend_context_uuid": frontend_context_uuid})

    return answer

def stream_search(query, mode=None, model=None, sources="web", language="en-US", thread_id=None):
    query = sanitize_query(query)
    sources_list = [s.strip() for s in sources.split(",")] if sources else ["web"]
    mode = mode or DEFAULT_MODE
    model_val = model if model and model in VALID_MODEL_NAMES else None
    thread = _get_thread(thread_id) if thread_id else None
    follow_up = None

    if thread and thread.get("backend_uuid"):
        follow_up = {"backend_uuid": thread["backend_uuid"], "frontend_context_uuid": thread.get("frontend_context_uuid", str(uuid.uuid4())), "attachments": []}

    client = _get_client(mode)
    backend_uuid = None
    frontend_context_uuid = None

    for chunk in client.search(query=query, mode=mode, model=model_val, sources=sources_list, language=language, stream=True, follow_up=follow_up):
        if "answer" in chunk:
            yield chunk["answer"]
        if not backend_uuid and chunk.get("backend_uuid"):
            backend_uuid = chunk["backend_uuid"]
        if not backend_uuid and chunk.get("uuid"):
            backend_uuid = chunk["uuid"]
        if not frontend_context_uuid and chunk.get("frontend_context_uuid"):
            frontend_context_uuid = chunk["frontend_context_uuid"]

    save_cookies(client.session.cookies.get_dict() if hasattr(client.session, "cookies") else {}, client.copilot, client.file_upload)
    if thread_id:
        _update_thread(thread_id, {"mode": mode, "backend_uuid": backend_uuid or (thread.get("backend_uuid") if thread else None), "frontend_context_uuid": frontend_context_uuid or str(uuid.uuid4())})


# ─── FastAPI App ─────────────────────────────────────────────────────────────

app = FastAPI(title="Web Search MCP", version="3.0.0")

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

# ─── SSE Transport for MCP ────────────────────────────────────────────────────
_sse_queue: asyncio.Queue | None = None

# ─── Shared MCP Handler ───────────────────────────────────────────────────────

async def _handle_mcp_jsonrpc(body: dict) -> dict | None:
    method = body.get("method", "")
    req_id = body.get("id")
    params = body.get("params", {})

    if method == "initialize":
        return {
            "jsonrpc": "2.0", "id": req_id,
            "result": {
                "protocolVersion": "2024-11-05",
                "capabilities": {"tools": {}},
                "serverInfo": {"name": "websearch", "version": "3.0.0"}
            }
        }

    if method == "notifications/initialized":
        return None

    if method == "ping":
        return {"jsonrpc": "2.0", "id": req_id, "result": {}}

    if method == "tools/list":
        return {"jsonrpc": "2.0", "id": req_id, "result": {"tools": MCP_TOOLS}}

    if method == "tools/call":
        tool_name = params.get("name", "ask")
        arguments = params.get("arguments", {})
        query = arguments.get("query", "")
        thread_id = arguments.get("thread_id") or str(uuid.uuid4())
        if not query:
            return {"jsonrpc": "2.0", "id": req_id, "error": {"code": -32602, "message": "Missing required parameter: query"}}
        mode = MCP_TOOL_MODE_MAP.get(tool_name, "pro")
        model = arguments.get("model") or (None if mode == "deep research" else (MCP_REASONING_DEFAULT_MODEL if mode == "reasoning" else MCP_DEFAULT_MODEL))
        loop = asyncio.get_event_loop()
        try:
            result = await loop.run_in_executor(EXECUTOR, lambda: search_sync(query, mode=mode, model=model, thread_id=thread_id))
            return {
                "jsonrpc": "2.0", "id": req_id,
                "result": {
                    "content": [{"type": "text", "text": result}],
                    "thread_id": thread_id,
                    "mode": mode,
                    "model": model
                }
            }
        except Exception as e:
            return {"jsonrpc": "2.0", "id": req_id, "error": {"code": -32603, "message": str(e)}}

    return {"jsonrpc": "2.0", "id": req_id, "error": {"code": -32601, "message": f"Method not found: {method}"}}


@app.post("/mcp")
async def mcp_endpoint(request: Request):
    body = await request.json()
    resp = await _handle_mcp_jsonrpc(body)
    if resp is None:
        return JSONResponse({}, status_code=200)
    return JSONResponse(resp)


@app.get("/sse")
async def sse_endpoint(request: Request):
    global _sse_queue
    queue: asyncio.Queue = asyncio.Queue()
    _sse_queue = queue

    async def event_generator():
        try:
            yield "event: endpoint\ndata: /messages\n\n"
            while True:
                data = await queue.get()
                if data is None:
                    break
                yield f"event: message\ndata: {json.dumps(data)}\n\n"
        except asyncio.CancelledError:
            pass
        finally:
            global _sse_queue
            if _sse_queue is queue:
                _sse_queue = None

    return StreamingResponse(event_generator(), media_type="text/event-stream")


@app.post("/messages")
async def messages_endpoint(request: Request):
    body = await request.json()
    resp = await _handle_mcp_jsonrpc(body)
    if resp is None:
        return JSONResponse({}, status_code=202)
    if _sse_queue is not None:
        await _sse_queue.put(resp)
        return JSONResponse({}, status_code=202)
    return JSONResponse(resp)


@app.get("/health")
async def health():
    stored = load_cookies()
    copilot = stored.get("copilot_remaining", -1) if stored else -1
    threads = _load_threads()
    uptime = round(time.time() - START_TIME, 0)
    account_age = round((time.time() - stored.get("saved_at", 0)) / 3600, 1) if stored else 0
    history = _load_history()
    total = len(history)
    ok = sum(1 for h in history if h.get("success"))
    avg_time = round(sum(h.get("elapsed", 0) for h in history) / total, 2) if total else 0
    return {
        "status": "ok", "version": "3.0.0",
        "uptime_hours": round(uptime / 3600, 1),
        "cookies_active": bool(stored and stored.get("cookies")),
        "copilot_remaining": copilot,
        "account_age_hours": account_age,
        "active_threads": len(threads),
        "total_queries": total,
        "success_rate": round(ok / total * 100, 1) if total else 0,
        "avg_response_time": avg_time
    }


@app.get("/status")
async def status_api():
    stored = load_cookies()
    history = _load_history()
    threads = _load_threads()
    uptime = round(time.time() - START_TIME)
    modes_used = {}
    for h in history:
        m = h.get("mode", "unknown")
        modes_used[m] = modes_used.get(m, 0) + 1
    return {
        "health": await health(),
        "recent_queries": history[-15:],
        "modes_used": modes_used,
        "thread_count": len(threads)
    }


# ─── MCP Protocol (JSON-RPC 2.0 over HTTP) ────────────────────────────────────

MCP_TOOLS = [
    {
        "name": "ask",
        "description": "Ask Perplexity a question using pro mode (web search + AI synthesis). Best for general questions.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "query": {"type": "string", "description": "The question or search query"},
                "model": {"type": "string", "description": "Model: gpt-5.4, claude-sonnet-4.6, claude-opus-4.7, gemini-3.1-pro, sonar, kimi-k2.6, nemotron-3-super-120b, gpt-5.5. All support -thinking suffix (default: gpt-5.4)"},
                "thread_id": {"type": "string", "description": "Optional thread ID for follow-up conversation"}
            },
            "required": ["query"]
        }
    },
    {
        "name": "search",
        "description": "Deep web search using Perplexity pro mode with comprehensive source analysis.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "query": {"type": "string", "description": "The search query"},
                "model": {"type": "string", "description": "Model: gpt-5.4, claude-sonnet-4.6, claude-opus-4.7, gemini-3.1-pro, sonar, kimi-k2.6, nemotron-3-super-120b, gpt-5.5. All support -thinking suffix (default: gpt-5.4)"},
                "thread_id": {"type": "string", "description": "Optional thread ID for follow-up"}
            },
            "required": ["query"]
        }
    },
    {
        "name": "reason",
        "description": "Extended reasoning and analysis using Perplexity reasoning mode. Best for complex problems requiring step-by-step thinking.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "query": {"type": "string", "description": "The question requiring deep reasoning"},
                "model": {"type": "string", "description": "Model: gpt-5.4-thinking, claude-sonnet-4.6-thinking, claude-opus-4.7-thinking, gemini-3.1-pro-thinking, kimi-k2.6-thinking, nemotron-3-super-120b-thinking, gpt-5.5-thinking (default: gpt-5.4-thinking)"},
                "thread_id": {"type": "string", "description": "Optional thread ID for follow-up"}
            },
            "required": ["query"]
        }
    },
    {
        "name": "research",
        "description": "Deep research using Perplexity deep research mode. Comprehensive multi-source analysis. Slower but thorough.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "query": {"type": "string", "description": "The research question or topic"},
                "thread_id": {"type": "string", "description": "Optional thread ID for follow-up"}
            },
            "required": ["query"]
        }
    }
]

MCP_TOOL_MODE_MAP = {"ask": "pro", "search": "pro", "reason": "reasoning", "research": "deep research"}
MCP_DEFAULT_MODEL = "gpt-5.4"
MCP_REASONING_DEFAULT_MODEL = "gpt-5.4-thinking"


# ─── OpenAI Compatible ───────────────────────────────────────────────────────

def _extract_query(messages):
    for m in reversed(messages):
        if m.get("role") == "user":
            return _content_to_str(m.get("content", ""))
    return ""

def _build_context_query(messages):
    if len(messages) <= 1:
        return _content_to_str(messages[-1].get("content", "")) if messages else ""
    context_parts = []
    for m in messages[:-1]:
        role = m.get("role", "user")
        content = _content_to_str(m.get("content", ""))
        if role == "system":
            continue
        prefix = "User" if role == "user" else "Assistant"
        context_parts.append(f"{prefix}: {content}")
    current_query = _content_to_str(messages[-1].get("content", "")) if messages else ""
    context = "\n".join(context_parts)
    if context:
        return f"[Previous conversation context]\n{context}\n\n[Current question]\n{current_query}"
    return current_query

def _build_chunk(req_id, created, model, delta, finish_reason=None):
    chunk = {"id": req_id, "object": "chat.completion.chunk", "created": created, "model": model, "choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}]}
    return json.dumps(chunk)


@app.post("/v1/chat/completions")
async def openai_chat_completions(request: Request):
    body = await request.json()
    messages = body.get("messages", [])
    model = body.get("model", "pro")
    stream = body.get("stream", False)
    language = body.get("language", "en-US")
    thread_id = body.get("thread_id")

    query = _extract_query(messages)
    if not query:
        return JSONResponse({"error": {"message": "No user message found", "type": "invalid_request_error"}}, status_code=400)

    if len(messages) > 2:
        query = _build_context_query(messages)

    mode = MODE_MAP.get(model, DEFAULT_MODE)
    model_val = model if model in VALID_MODEL_NAMES else None

    if not thread_id:
        thread_id = str(uuid.uuid4())
    if not _get_thread(thread_id):
        _update_thread(thread_id, {"mode": mode})

    req_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
    created = int(time.time())

    if stream:
        async def generate():
            full_text = ""
            try:
                for chunk_text in stream_search(query, mode, model_val, "web", language, thread_id=thread_id):
                    if chunk_text and chunk_text.startswith("Error:"):
                        delta_text = chunk_text
                    elif chunk_text:
                        delta_text = chunk_text[len(full_text):]
                        full_text = chunk_text
                    else:
                        continue
                    if delta_text:
                        yield f"data: {_build_chunk(req_id, created, model, {'content': delta_text})}\n\n"
                yield f"data: {_build_chunk(req_id, created, model, {}, 'stop')}\n\n"
                yield "data: [DONE]\n\n"
                _update_thread(thread_id, {"messages": messages + [{"role": "assistant", "content": full_text}]})
            except Exception as e:
                yield f"data: {_build_chunk(req_id, created, model, {'content': f'Error: {str(e)}'}, 'stop')}\n\n"
                yield "data: [DONE]\n\n"
        return StreamingResponse(generate(), media_type="text/event-stream")
    else:
        try:
            result = search_sync(query, mode, model_val, "web", language, thread_id=thread_id)
            _update_thread(thread_id, {"messages": messages + [{"role": "assistant", "content": result}]})
            return JSONResponse({
                "id": req_id, "object": "chat.completion", "created": created, "model": model,
                "choices": [{"index": 0, "message": {"role": "assistant", "content": result}, "finish_reason": "stop"}],
                "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
                "thread_id": thread_id,
            })
        except Exception as e:
            return JSONResponse({"error": {"message": str(e), "type": "server_error"}}, status_code=500)


@app.get("/v1/models")
async def openai_list_models():
    models = [{"id": "auto", "object": "model", "owned_by": "perplexity"}, {"id": "pro", "object": "model", "owned_by": "perplexity"}, {"id": "reasoning", "object": "model", "owned_by": "perplexity"}, {"id": "deep-research", "object": "model", "owned_by": "perplexity"}]
    for mode_models in MODEL_MAPPINGS.values():
        for k in mode_models:
            if k is not None:
                models.append({"id": k, "object": "model", "owned_by": "perplexity"})
    for lab in LABS_MODELS:
        models.append({"id": lab, "object": "model", "owned_by": "perplexity-labs"})
    return JSONResponse({"object": "list", "data": models})


# ─── Web UI ───────────────────────────────────────────────────────────────────

UI_HTML = """<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0, maximum-scale=1.0, user-scalable=no">
<title>Web Search MCP</title>
<style>
*{margin:0;padding:0;box-sizing:border-box}
:root{--bg:#0a0a0f;--surface:#12121a;--surface2:#1a1a25;--border:#2a2a3a;--text:#e0e0e0;--text2:#888;--accent:#6c5ce7;--accent2:#a29bfe;--green:#00b894;--red:#e74c3c;--orange:#fdcb6e}
html,body{height:100%;overflow:hidden}
body{font-family:-apple-system,BlinkMacSystemFont,'Segoe UI',system-ui,sans-serif;background:var(--bg);color:var(--text);display:flex;flex-direction:column;height:100vh;height:100dvh}
.header{background:var(--surface);border-bottom:1px solid var(--border);padding:10px 16px;display:flex;align-items:center;justify-content:space-between;flex-shrink:0}
.header h1{font-size:16px;background:linear-gradient(135deg,var(--accent),var(--accent2));-webkit-background-clip:text;-webkit-text-fill-color:transparent;font-weight:700}
.header .status{display:flex;gap:8px;align-items:center;font-size:11px;color:var(--text2);flex-shrink:0;max-width:60%;overflow:hidden;text-overflow:ellipsis;white-space:nowrap}
.header .status .dot{width:6px;height:6px;border-radius:50%;background:var(--green);animation:pulse 2s infinite;flex-shrink:0}
@keyframes pulse{0%,100%{opacity:1}50%{opacity:.4}}
.tabs{display:flex;background:var(--surface);border-bottom:1px solid var(--border);flex-shrink:0}
.tab{padding:8px 16px;cursor:pointer;font-size:12px;color:var(--text2);border-bottom:2px solid transparent;transition:all .2s;flex:1;text-align:center}
.tab:hover{color:var(--text)}.tab.active{color:var(--accent2);border-bottom-color:var(--accent)}
.content{flex:1;overflow:hidden;display:flex;flex-direction:column;min-height:0}
.panel{display:none;flex:1;flex-direction:column;overflow:hidden;min-height:0}.panel.active{display:flex}
.chat-area{flex:1;display:flex;flex-direction:column;overflow:hidden;min-height:0}
.messages{flex:1;display:flex;flex-direction:column;gap:10px;overflow-y:auto;padding:16px;-webkit-overflow-scrolling:touch;overscroll-behavior:contain}
.messages:empty::after{content:'Ask me anything...';color:var(--text2);text-align:center;margin:auto;font-size:14px}
.msg{max-width:85%;padding:10px 14px;border-radius:12px;font-size:13px;line-height:1.5;white-space:pre-wrap;word-break:break-word;overflow-wrap:break-word;position:relative}
@media(min-width:600px){.msg{max-width:70%}}
.msg.user{align-self:flex-end;background:var(--accent);color:#fff;border-bottom-right-radius:4px}
.msg.assistant{align-self:flex-start;background:var(--surface2);border:1px solid var(--border);border-bottom-left-radius:4px}
.msg.loading{color:var(--text2);font-style:italic}
.chat-input{display:flex;gap:6px;padding:10px 12px;background:var(--surface);border-top:1px solid var(--border);flex-shrink:0;align-items:stretch}
.chat-input select{background:var(--surface2);border:1px solid var(--border);color:var(--text);padding:8px;border-radius:8px;font-family:inherit;font-size:12px;flex-shrink:0;min-width:0}
.chat-input input{flex:1;min-width:0;background:var(--surface2);border:1px solid var(--border);color:var(--text);padding:8px 12px;border-radius:8px;font-family:inherit;font-size:13px;outline:none}
.chat-input input:focus{border-color:var(--accent)}
.chat-input button{background:var(--accent);color:#fff;border:none;padding:8px 14px;border-radius:8px;cursor:pointer;font-family:inherit;font-size:13px;font-weight:600;flex-shrink:0}
.chat-input button:hover{background:var(--accent2)}
.chat-input button:disabled{opacity:.5;cursor:not-allowed}
.mcp-container{flex:1;display:flex;flex-direction:column;padding:16px;gap:12px;overflow-y:auto;-webkit-overflow-scrolling:touch}
.tool-card{background:var(--surface2);border:1px solid var(--border);border-radius:10px;padding:14px;cursor:pointer;transition:all .2s}
.tool-card:hover{border-color:var(--accent)}
.tool-card h3{font-size:13px;color:var(--accent2);margin-bottom:4px}
.tool-card p{font-size:11px;color:var(--text2);margin-bottom:8px;line-height:1.4}
.tool-test{display:none;margin-top:10px;gap:8px;flex-direction:column}
.tool-card.open .tool-test{display:flex}
.tool-test input{background:var(--bg);border:1px solid var(--border);color:var(--text);padding:10px;border-radius:8px;font-family:inherit;font-size:12px;width:100%}
.tool-test .tool-actions{display:flex;gap:8px}
.tool-test button{background:var(--accent);color:#fff;border:none;padding:8px 14px;border-radius:8px;cursor:pointer;font-family:inherit;font-size:12px;flex-shrink:0}
.tool-test button:hover{background:var(--accent2)}
.tool-result{background:var(--bg);border:1px solid var(--border);border-radius:8px;padding:10px;font-size:12px;white-space:pre-wrap;max-height:250px;overflow-y:auto;color:var(--green);line-height:1.5;word-break:break-word}
.status-wrap{flex:1;overflow-y:auto;padding:16px;display:flex;flex-direction:column;gap:16px;-webkit-overflow-scrolling:touch}
.s-row{display:grid;grid-template-columns:repeat(4,1fr);gap:8px}
@media(max-width:500px){.s-row{grid-template-columns:repeat(2,1fr)}}
.s-card{background:var(--surface2);border:1px solid var(--border);border-radius:8px;padding:10px 12px;display:flex;flex-direction:column}
.s-card .label{font-size:9px;color:var(--text2);text-transform:uppercase;letter-spacing:.5px;margin-bottom:4px}
.s-card .val{font-size:18px;font-weight:700;color:var(--accent2)}
.s-card .val.ok{color:var(--green)}.s-card .val.bad{color:var(--red)}.s-card .val.warn{color:var(--orange)}
.s-card .unit{font-size:10px;color:var(--text2);margin-top:2px}
.s-section{background:var(--surface2);border:1px solid var(--border);border-radius:8px;padding:12px}
.s-section h3{font-size:11px;color:var(--text2);text-transform:uppercase;letter-spacing:.5px;margin-bottom:8px;padding-bottom:6px;border-bottom:1px solid var(--border)}
.s-table{width:100%;font-size:11px;border-collapse:collapse}
.s-table th{text-align:left;color:var(--text2);font-weight:400;padding:4px 6px;border-bottom:1px solid var(--border)}
.s-table td{padding:5px 6px;border-bottom:1px solid var(--border);vertical-align:top;max-width:200px;overflow:hidden;text-overflow:ellipsis;white-space:nowrap}
.s-table tr:last-child td{border-bottom:none}
.s-table .ok{color:var(--green)}.s-table .fail{color:var(--red)}
.s-table .time{color:var(--text2);font-size:10px;white-space:nowrap}
.model-chips{display:flex;flex-wrap:wrap;gap:4px}
.model-chip{background:var(--bg);border:1px solid var(--border);padding:3px 8px;border-radius:4px;font-size:10px;color:var(--text)}
.model-chip.mode{border-color:var(--accent);color:var(--accent2)}
.api-ref{font-size:11px;line-height:1.8}
.api-ref code{background:var(--bg);padding:2px 6px;border-radius:4px;font-size:10px;color:var(--green)}
.api-ref .method{color:var(--accent2);font-weight:600;display:inline-block;min-width:35px}
.loading{color:var(--orange);font-size:12px}
</style>
</head>
<body>
<div class="header">
  <h1>Web Search</h1>
  <div class="status"><div class="dot"></div><span id="headerStatus">...</span></div>
</div>
<div class="tabs">
  <div class="tab active" onclick="showTab('chat')">Chat</div>
  <div class="tab" onclick="showTab('mcp')">MCP Tools</div>
  <div class="tab" onclick="showTab('status')">Status</div>
</div>
<div class="content">
  <div class="panel active" id="panel-chat">
    <div class="chat-area">
      <div class="messages" id="messages"></div>
      <div class="chat-input">
        <select id="chatMode"><option value="pro">Pro</option><option value="reasoning">Reason</option><option value="deep-research">Research</option></select>
        <input id="chatInput" placeholder="Ask anything..." onkeydown="if(event.key==='Enter'&&!event.shiftKey){event.preventDefault();sendChat()}">
        <button id="sendBtn" onclick="sendChat()">Send</button>
      </div>
    </div>
  </div>
  <div class="panel" id="panel-mcp">
    <div class="mcp-container" id="mcpTools"></div>
  </div>
  <div class="panel" id="panel-status">
    <div class="status-wrap" id="statusWrap"></div>
  </div>
</div>
<script>
let threadId=null,sending=false;
const $=id=>document.getElementById(id);
function esc(s){if(!s)return'';return s.replace(/&/g,'&amp;').replace(/</g,'&lt;').replace(/>/g,'&gt;')}
async function init(){
  try{
    const h=await fetch('/health').then(r=>r.json());
    $('headerStatus').textContent='v'+h.version+' | '+(h.cookies_active?'cookies ok':'no cookies')+' | copilot: '+h.copilot_remaining;
  }catch(e){$('headerStatus').textContent='error'}
  loadMcpTools();loadStatus();
}
function showTab(t){['chat','mcp','status'].forEach((n,i)=>{document.querySelectorAll('.tab')[i].classList.toggle('active',n===t);document.querySelectorAll('.panel')[i].classList.toggle('active',n===t)});if(t==='status')loadStatus()}
function scrollBottom(){const m=$('messages');m.scrollTop=m.scrollHeight}
function addMsg(role,text){
  const d=document.createElement('div');d.className='msg '+role+(text==='...'?' loading':'');d.textContent=text;
  $('messages').appendChild(d);scrollBottom();return d;
}
function setSending(v){sending=v;$('sendBtn').disabled=v;$('chatInput').disabled=v}
async function sendChat(){
  if(sending)return;
  const q=$('chatInput').value.trim();if(!q)return;
  $('chatInput').value='';setSending(true);
  const userMsg=addMsg('user',q);const botMsg=addMsg('assistant','...');
  try{
    const body={model:$('chatMode').value,messages:[{role:'user',content:q}],stream:false};
    if(threadId)body.thread_id=threadId;
    const r=await fetch('/v1/chat/completions',{method:'POST',headers:{'Content-Type':'application/json'},body:JSON.stringify(body)});
    const d=await r.json();
    if(d.error){botMsg.textContent='Error: '+d.error.message;botMsg.classList.add('loading')}
    else{botMsg.textContent=d.choices[0].message.content;botMsg.classList.remove('loading');threadId=d.thread_id}
  }catch(e){botMsg.textContent='Error: '+e.message;botMsg.classList.add('loading')}
  setSending(false);scrollBottom();
}
async function loadMcpTools(){
  try{
    const r=await fetch('/mcp',{method:'POST',headers:{'Content-Type':'application/json'},body:JSON.stringify({jsonrpc:'2.0',id:1,method:'tools/list'})});
    const d=await r.json();const el=$('mcpTools');el.innerHTML='';
    d.result.tools.forEach(t=>{
      const card=document.createElement('div');card.className='tool-card';
      card.innerHTML='<h3>'+t.name+'</h3><p>'+t.description+'</p><div class="tool-test"><input id="input-'+t.name+'" placeholder="Enter query..." onkeydown="if(event.key===\\'Enter\\')callTool(\\''+t.name+'\\')"><div class="tool-actions"><button onclick="callTool(\\''+t.name+'\\')">Execute</button></div><div id="result-'+t.name+'" class="tool-result" style="display:none"></div></div>';
      card.onclick=function(e){if(e.target.tagName==='INPUT'||e.target.tagName==='BUTTON')return;card.classList.toggle('open')};
      el.appendChild(card);
    });
  }catch(e){$('mcpTools').innerHTML='<p class="loading">Failed to load tools</p>'}
}
async function callTool(name){
  const q=$('input-'+name).value.trim();if(!q)return;
  const res=$('result-'+name);res.style.display='block';res.textContent='Loading...';
  try{
    const r=await fetch('/mcp',{method:'POST',headers:{'Content-Type':'application/json'},body:JSON.stringify({jsonrpc:'2.0',id:2,method:'tools/call',params:{name:name,arguments:{query:q}}})});
    const d=await r.json();
    if(d.error)res.textContent='Error: '+d.error.message;
    else res.textContent=d.result.content.map(function(c){return c.text}).join('\\n');
  }catch(e){res.textContent='Error: '+e.message}
}
async function loadStatus(){
  try{
    const[s,mr]=await Promise.all([fetch('/status').then(r=>r.json()),fetch('/v1/models').then(r=>r.json())]);
    const h=s.health;const models=mr.data;const recent=s.recent_queries;const modes=s.modes_used;
    const el=$('statusWrap');
    function fmtTime(iso){if(!iso)return'-';const d=new Date(iso);return d.toLocaleTimeString([],{hour:'2-digit',minute:'2-digit',second:'2-digit'})}
    function fmtDur(s){if(s<60)return s+'s';if(s<3600)return Math.round(s/60)+'m';return (s/3600).toFixed(1)+'h'}
    let row1='<div class="s-row">'
      +'<div class="s-card"><span class="label">Account</span><span class="val '+(h.cookies_active?'ok':'bad')+'">'+(h.cookies_active?'Active':'Dead')+'</span><span class="unit">'+(h.account_age_hours?h.account_age_hours+'h old':'')+'</span></div>'
      +'<div class="s-card"><span class="label">Copilot</span><span class="val '+(h.copilot_remaining>0?'ok':'bad')+'">'+h.copilot_remaining+'</span><span class="unit">credits left</span></div>'
      +'<div class="s-card"><span class="label">Queries</span><span class="val">'+h.total_queries+'</span><span class="unit">'+h.success_rate+'% ok | ~'+h.avg_response_time+'s avg</span></div>'
      +'<div class="s-card"><span class="label">Uptime</span><span class="val">'+fmtDur(h.uptime_hours*3600)+'</span><span class="unit">'+h.active_threads+' threads</span></div>'
      +'</div>';
    let modesHtml='<div class="s-section"><h3>Modes Used</h3><div style="display:flex;gap:12px;flex-wrap:wrap;font-size:12px">';
    Object.entries(modes).forEach(function(e){modesHtml+='<span style="color:var(--accent2)">'+e[0]+'</span> <span style="color:var(--text2)">'+e[1]+'x</span>'});
    if(!Object.keys(modes).length)modesHtml+='<span style="color:var(--text2)">No queries yet</span>';
    modesHtml+='</div></div>';
    let recentHtml='';
    if(recent.length){
      recentHtml='<div class="s-section"><h3>Recent Queries</h3><table class="s-table"><thead><tr><th>Time</th><th>Query</th><th>Mode</th><th>Took</th><th></th></tr></thead><tbody>';
      recent.slice().reverse().forEach(function(q){
        recentHtml+='<tr><td class="time">'+fmtTime(q.ts)+'</td><td>'+esc(q.query)+'</td><td>'+q.mode+'</td><td>'+q.elapsed+'s</td><td class="'+(q.success?'ok':'fail')+'">'+(q.success?'ok':'fail')+'</td></tr>';
      });
      recentHtml+='</tbody></table></div>';
    }
    let modeNames=['auto','pro','reasoning','deep-research','sonar'];
    let modelList=models.filter(function(m){return modeNames.indexOf(m.id)===-1}).map(function(m){return m.id});
    let modeChips=models.filter(function(m){return modeNames.indexOf(m.id)!==-1}).map(function(m){return '<span class="model-chip mode">'+m.id+'</span>'}).join('');
    let modelChips=modelList.map(function(m){return '<span class="model-chip">'+m+'</span>'}).join('');
    let modelsHtml='<div class="s-section"><h3>Modes</h3><div class="model-chips" style="margin-bottom:10px">'+modeChips+'</div><h3>Models</h3><div class="model-chips">'+modelChips+'</div></div>';
    let apiHtml='<div class="s-section"><h3>API Reference</h3><div class="api-ref">'
      +'<div><span class="method">POST</span> <code>/v1/chat/completions</code> — OpenAI-compatible chat</div>'
      +'<div><span class="method">POST</span> <code>/mcp</code> — JSON-RPC 2.0 (tools/list, tools/call)</div>'
      +'<div><span class="method">GET</span> <code>/v1/models</code> — List available models</div>'
      +'<div><span class="method">GET</span> <code>/health</code> — Health check</div>'
      +'<div><span class="method">GET</span> <code>/status</code> — Full status + history</div>'
      +'</div></div>';
    el.innerHTML=row1+modesHtml+recentHtml+modelsHtml+apiHtml;
  }catch(e){$('statusWrap').innerHTML='<p class="loading">Failed to load: '+e.message+'</p>'}
}
init();
</script>
</body>
</html>"""


@app.get("/", response_class=HTMLResponse)
async def web_ui():
    return UI_HTML


if __name__ == "__main__":
    import uvicorn
    uvicorn.run("app:app", host="0.0.0.0", port=7860, workers=1, timeout_keep_alive=300, log_level="info")