import os
import re
import json
import time
import uuid
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, LabsClient, Emailnator
from perplexity.config import DEFAULT_HEADERS, EMAILNATOR_BASE_URL, EMAILNATOR_HEADERS, MODEL_MAPPINGS, LABS_MODELS
from perplexity.exceptions import PerplexityError, AuthenticationError, QueryLimitExceededError, AccountCreationError
from perplexity.utils import retry_with_backoff, sanitize_query
from perplexity.logger import get_logger

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

logger = get_logger("app")

ADMIN_PASSWORD = os.environ.get("ADMIN_PASSWORD", "")
API_KEY = os.environ.get("API_KEY", "")

_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()) / "pplx"
DATA_DIR.mkdir(parents=True, exist_ok=True)

COOKIES_FILE = DATA_DIR / "pplx_cookies.json"
EMAILNATOR_COOKIES_FILE = DATA_DIR / "emailnator_cookies.json"
HISTORY_FILE = DATA_DIR / "pplx_history.json"
LOG_FILE = DATA_DIR / "perplexity.log"
THREADS_FILE = DATA_DIR / "pplx_threads.json"

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)

MODE_MAP = {
    "auto": "auto",
    "pro": "pro",
    "reasoning": "reasoning",
    "deep-research": "deep research",
    "deep_research": "deep research",
    "sonar": "pro",
    "sonar2": "pro",
    "sonar2-pro": "pro",
    "sonar2-deep-research": "deep research",
}

VALID_MODEL_NAMES = set()
for mode_models in MODEL_MAPPINGS.values():
    for k in mode_models:
        if k is not None:
            VALID_MODEL_NAMES.add(k)


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)


# ─── Thread / Conversation 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):
    threads = _load_threads()
    return 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": "auto",
        }
    threads[thread_id].update(data)
    _save_threads(threads)
    return threads[thread_id]


def _build_context_query(messages):
    if len(messages) <= 1:
        return None, _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:
        full_query = f"[Previous conversation context]\n{context}\n\n[Current question]\n{current_query}"
    else:
        full_query = current_query
    return context, full_query


# ─── 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):
    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_for_mode(mode):
    needs_auth = mode in ("pro", "reasoning", "deep research")
    if not needs_auth:
        return Client()
    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:
            logger.info("Copilot credits exhausted, creating new account...")
        else:
            return client
    logger.info("Creating new account...")
    future = EXECUTOR.submit(auto_create_account)
    try:
        cookies = future.result(timeout=ACCOUNT_CREATE_TIMEOUT)
    except FuturesTimeoutError:
        raise AccountCreationError("Account creation timed out after {}s".format(ACCOUNT_CREATE_TIMEOUT))
    client = Client(cookies)
    save_cookies(cookies, client.copilot, client.file_upload)
    return client


# ─── Search with Thread Context ────────────────────────────────────────────────

def search_sync(query, mode="auto", model=None, sources="web", language="en-US", cookies_str="", follow_up=None, thread_id=None):
    try:
        query = sanitize_query(query)
        sources_list = [s.strip() for s in sources.split(",")] if sources else ["web"]
        model_val = model if model and model in VALID_MODEL_NAMES else None
        thread = _get_thread(thread_id) if thread_id else 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}: backend_uuid={thread['backend_uuid']}")

        if cookies_str and cookies_str.strip():
            try:
                manual_cookies = json.loads(cookies_str)
                client = Client(manual_cookies)
                save_cookies(manual_cookies, client.copilot, client.file_upload)
            except Exception:
                client = _get_client_for_mode(mode)
        else:
            client = _get_client_for_mode(mode)

        resp = client.search(
            query=query,
            mode=mode,
            model=model_val,
            sources=sources_list,
            language=language,
            follow_up=follow_up,
        )

        if mode in ("pro", "reasoning", "deep research"):
            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")

        if thread_id:
            backend_uuid = resp.get("backend_uuid") or resp.get("uuid") or (resp.get("text", [{}])[-1].get("uuid") if isinstance(resp.get("text"), list) else None)
            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
    except PerplexityError as e:
        logger.error(f"Search error: {e}")
        if thread_id and thread and thread.get("backend_uuid"):
            logger.info(f"Session may be expired for thread {thread_id}, retrying without follow_up")
            thread["backend_uuid"] = None
            _update_thread(thread_id, {"backend_uuid": None})
        raise
    except Exception as e:
        logger.error(f"Unexpected search error: {e}")
        raise


def stream_search(query, mode="auto", model=None, sources="web", language="en-US", cookies_str="", follow_up=None, thread_id=None):
    try:
        query = sanitize_query(query)
        sources_list = [s.strip() for s in sources.split(",")] if sources else ["web"]
        model_val = model if model and model in VALID_MODEL_NAMES else None
        thread = _get_thread(thread_id) if thread_id else 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}: backend_uuid={thread['backend_uuid']}")

        if cookies_str and cookies_str.strip():
            try:
                manual_cookies = json.loads(cookies_str)
                client = Client(manual_cookies)
                save_cookies(manual_cookies, client.copilot, client.file_upload)
            except Exception:
                client = _get_client_for_mode(mode)
        else:
            client = _get_client_for_mode(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"]

        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()),
            })

    except PerplexityError as e:
        logger.error(f"Stream search error: {e}")
        if thread_id and thread and thread.get("backend_uuid"):
            logger.info(f"Session may be expired for thread {thread_id}, clearing backend_uuid")
            _update_thread(thread_id, {"backend_uuid": None})
        yield "Error: {}".format(str(e))
    except Exception as e:
        logger.error(f"Unexpected stream search error: {e}")
        yield "Error: {}".format(str(e))


# ─── Logging ───────────────────────────────────────────────────────────────────

def _log_query(query, mode, model, answer_len, success=True, thread_id=None):
    try:
        history = []
        if HISTORY_FILE.exists():
            with open(HISTORY_FILE, "r") as f:
                history = json.load(f)
        history.append({
            "timestamp": datetime.utcnow().isoformat(),
            "query": query[:200],
            "mode": mode,
            "model": model or "default",
            "answer_len": answer_len,
            "success": success,
            "thread_id": thread_id,
        })
        history = history[-500:]
        with open(HISTORY_FILE, "w") as f:
            json.dump(history, f)
    except Exception:
        pass


# ─── OpenAI Format Helpers ─────────────────────────────────────────────────────

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


def _extract_sources(messages):
    for m in messages:
        if m.get("role") == "system":
            s = _content_to_str(m.get("content", "")).lower()
            if "scholar" in s:
                return "scholar"
            if "news" in s:
                return "news"
    return "web"


def _extract_thread_id(messages):
    for m in messages:
        if m.get("role") == "system":
            tid = re.search(r"thread[_\-:]?\s*(\S+)", _content_to_str(m.get("content", "")), re.I)
            if tid:
                return tid.group(1)
    return None


def _build_openai_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)


def _build_openai_response(req_id, created, model, content):
    return {
        "id": req_id,
        "object": "chat.completion",
        "created": created,
        "model": model,
        "choices": [{"index": 0, "message": {"role": "assistant", "content": content}, "finish_reason": "stop"}],
        "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
    }


def _extract_query_from_input(input_data):
    if isinstance(input_data, str):
        return input_data
    if isinstance(input_data, list):
        for item in reversed(input_data):
            if isinstance(item, dict) and item.get("role") == "user":
                return _content_to_str(item.get("content", ""))
            if isinstance(item, dict) and item.get("type") == "message":
                return _content_to_str(item.get("content", ""))
    return ""


def _extract_system_from_input(input_data):
    if isinstance(input_data, list):
        for item in input_data:
            if isinstance(item, dict):
                role = item.get("role", "")
                if role == "system":
                    return _content_to_str(item.get("content", ""))
                if item.get("type") == "system":
                    return _content_to_str(item.get("content", ""))
    return ""


def _build_responses_sse_event(event_type, data_dict):
    return "event: {}\ndata: {}\n\n".format(event_type, json.dumps(data_dict))


# ─── Admin UI ──────────────────────────────────────────────────────────────────

ADMIN_HTML = """<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>PPLX v2 Admin</title>
<style>
*{margin:0;padding:0;box-sizing:border-box}
body{font-family:-apple-system,BlinkMacSystemFont,'Segoe UI',Roboto,sans-serif;background:#0f172a;color:#e2e8f0;min-height:100vh}
.header{background:linear-gradient(135deg,#1e293b 0%,#0f172a 100%);padding:1.5rem 2rem;border-bottom:1px solid #334155;display:flex;align-items:center;justify-content:space-between}
.header h1{font-size:1.5rem;font-weight:700;background:linear-gradient(135deg,#38bdf8,#818cf8);-webkit-background-clip:text;-webkit-text-fill-color:transparent}
.header .badge{background:#1e40af;color:#93c5fd;padding:0.25rem 0.75rem;border-radius:9999px;font-size:0.75rem;font-weight:600}
.container{max-width:1200px;margin:2rem auto;padding:0 2rem}
.grid{display:grid;grid-template-columns:repeat(auto-fit,minmax(300px,1fr));gap:1.5rem}
.card{background:#1e293b;border:1px solid #334155;border-radius:1rem;padding:1.5rem}
.card h2{font-size:1rem;font-weight:600;color:#94a3b8;margin-bottom:1rem;text-transform:uppercase;letter-spacing:0.05em}
.stat{font-size:2.5rem;font-weight:700;line-height:1}
.stat.ok{color:#34d399}.stat.warn{color:#fbbf24}.stat.err{color:#f87171}
.stat-label{font-size:0.8rem;color:#64748b;margin-top:0.25rem}
.btn{padding:0.625rem 1.25rem;border-radius:0.5rem;font-size:0.875rem;font-weight:600;border:none;cursor:pointer;transition:all 0.15s}
.btn-primary{background:#2563eb;color:white}.btn-primary:hover{background:#1d4ed8}
.btn-danger{background:#dc2626;color:white}.btn-danger:hover{background:#b91c1c}
.btn-success{background:#059669;color:white}.btn-success:hover{background:#047857}
.btn-group{display:flex;gap:0.5rem;flex-wrap:wrap;margin-top:1rem}
.log{font-family:'SF Mono',Monaco,'Cascadia Code',monospace;font-size:0.75rem;background:#0f172a;border:1px solid #334155;border-radius:0.5rem;padding:1rem;max-height:300px;overflow-y:auto;line-height:1.6}
.log .ok{color:#34d399}.log .err{color:#f87171}.log .warn{color:#fbbf24}.log .info{color:#38bdf8}
table{width:100%;border-collapse:collapse;font-size:0.8rem}
th{text-align:left;padding:0.5rem;color:#64748b;border-bottom:1px solid #334155;font-weight:600}
td{padding:0.5rem;border-bottom:1px solid #1e293b}
.mono{font-family:'SF Mono',Monaco,monospace;font-size:0.75rem}
.tabs{display:flex;gap:0;border-bottom:2px solid #334155;margin-bottom:1.5rem}
.tab{padding:0.75rem 1.5rem;cursor:pointer;color:#94a3b8;font-weight:600;font-size:0.875rem;border-bottom:2px solid transparent;margin-bottom:-2px;transition:all 0.15s}
.tab:hover{color:#e2e8f0}.tab.active{color:#38bdf8;border-bottom-color:#38bdf8}
.panel{display:none}.panel.active{display:block}
.status-dot{width:8px;height:8px;border-radius:50%;display:inline-block;margin-right:0.5rem}
.status-dot.ok{background:#34d399}.status-dot.err{background:#f87171}.status-dot.warn{background:#fbbf24}
.v2-tag{background:linear-gradient(135deg,#818cf8,#c084fc);color:white;padding:0.15rem 0.5rem;border-radius:9999px;font-size:0.65rem;font-weight:700;margin-left:0.5rem}
</style>
</head>
<body>
<div class="header">
<h1>PPLX Admin<span class="v2-tag">v2</span></h1>
<span class="badge">Multi-Turn Proxy</span>
</div>
<div class="container">
<div class="tabs">
<div class="tab active" onclick="switchTab('status')">Status</div>
<div class="tab" onclick="switchTab('accounts')">Accounts</div>
<div class="tab" onclick="switchTab('threads')">Threads</div>
<div class="tab" onclick="switchTab('history')">History</div>
<div class="tab" onclick="switchTab('models')">Models</div>
<div class="tab" onclick="switchTab('logs')">Logs</div>
</div>
<div id="status" class="panel active">
<div class="grid">
<div class="card"><h2>Copilot Credits</h2><div class="stat" id="copilot-count">--</div><div class="stat-label">remaining pro queries</div></div>
<div class="card"><h2>Session Age</h2><div class="stat" id="session-age">--</div><div class="stat-label">since last account creation</div></div>
<div class="card"><h2>File Uploads</h2><div class="stat" id="file-uploads">--</div><div class="stat-label">remaining uploads</div></div>
<div class="card"><h2>Service</h2><div style="display:flex;align-items:center;margin-bottom:0.5rem"><span class="status-dot" id="service-dot"></span><span id="service-status">Checking...</span></div><div class="stat-label" id="uptime-label"></div></div>
</div>
<div class="card" style="margin-top:1.5rem"><h2>Quick Actions</h2><div class="btn-group">
<button class="btn btn-primary" onclick="createAccount()">Create New Account</button>
<button class="btn btn-danger" onclick="clearSession()">Clear Session</button>
<button class="btn btn-success" onclick="refreshStatus()">Refresh Status</button>
</div><div id="action-result" class="log" style="margin-top:1rem;display:none"></div></div>
</div>
<div id="accounts" class="panel">
<div class="card"><h2>Account Management</h2>
<p style="color:#64748b;margin-bottom:1rem;font-size:0.875rem">Auto-creates free Perplexity accounts. Each account gets 5 copilot queries and 10 file uploads. v2 preserves sessions for multi-turn conversations.</p>
<div class="btn-group"><button class="btn btn-primary" onclick="createAccount()">Create New Account</button><button class="btn btn-success" onclick="checkAccount()">Check Current Account</button></div>
<div id="account-result" class="log" style="margin-top:1rem;display:none"></div></div>
</div>
<div id="threads" class="panel">
<div class="card"><h2>Conversation Threads</h2>
<p style="color:#64748b;margin-bottom:1rem;font-size:0.875rem">Active threads with Perplexity session context. Backend UUID enables follow-up queries without resending full history.</p>
<div id="threads-table" style="overflow-x:auto"></div></div>
</div>
<div id="history" class="panel">
<div class="card"><h2>Query History</h2><div id="history-table" style="overflow-x:auto"></div></div>
</div>
<div id="models" class="panel">
<div class="grid">
<div class="card"><h2>Search Modes & Models</h2><table><thead><tr><th>Mode</th><th>Model</th><th>Internal</th><th>Auth</th></tr></thead><tbody id="models-table"></tbody></table></div>
<div class="card"><h2>Labs Models (WebSocket)</h2><table><thead><tr><th>Model</th><th>Protocol</th></tr></thead><tbody id="labs-table"></tbody></table></div>
</div>
</div>
<div id="logs" class="panel">
<div class="card"><h2>Server Logs</h2><div class="log" id="log-output">Loading...</div></div>
</div>
</div>
<script>
const API='/admin/api';
const ADMIN_TOKEN=localStorage.getItem('admin_token')||'';
function switchTab(id){
document.querySelectorAll('.tab').forEach(t=>t.classList.remove('active'));
document.querySelectorAll('.panel').forEach(p=>p.classList.remove('active'));
event.target.classList.add('active');
document.getElementById(id).classList.add('active');
if(id==='threads')loadThreads();
}
async function api(path,opts={}){
try{
const headers=opts.headers||{};
if(ADMIN_TOKEN)headers['Authorization']='Bearer '+ADMIN_TOKEN;
opts.headers=headers;
const r=await fetch(API+path,opts);return await r.json()
}catch(e){return{error:e.message}}
}
async function refreshStatus(){
const s=await api('/status');
const el=(id,v)=>{const e=document.getElementById(id);if(e)e.textContent=v};
if(s.copilot_remaining!==undefined){
if(s.copilot_remaining===-1){el('copilot-count','\\u221e');var cc=document.getElementById('copilot-count');cc.className='stat ok'}
else{el('copilot-count',s.copilot_remaining);var cc=document.getElementById('copilot-count');cc.className='stat '+(s.copilot_remaining>2?'ok':s.copilot_remaining>0?'warn':'err')}
}
if(s.file_upload_remaining!==undefined)el('file-uploads',s.file_upload_remaining);
if(s.session_age!==undefined){const h=Math.floor(s.session_age/3600);const m=Math.floor((s.session_age%3600)/60);el('session-age',h>0?h+'h '+m+'m':m+'m');}
const dot=document.getElementById('service-dot');const st=document.getElementById('service-status');
if(s.cookies_active){dot.className='status-dot ok';st.textContent='Active'}
else{dot.className='status-dot warn';st.textContent='No session (auto mode only)'}
}
async function createAccount(){showResult('action-result','Creating account...','info');const r=await api('/account/create',{method:'POST'});if(r.ok)showResult('action-result','Account created! Copilot: '+r.copilot_remaining+', Uploads: '+r.file_upload_remaining,'ok');else showResult('action-result','Error: '+r.error,'err');refreshStatus();}
async function clearSession(){const r=await api('/session/clear',{method:'POST'});showResult('action-result',r.ok?'Session cleared':'Error: '+r.error,r.ok?'ok':'err');refreshStatus();}
async function checkAccount(){const r=await api('/account/check');if(r.ok)showResult('account-result','Copilot: '+r.copilot_remaining+' | Uploads: '+r.file_upload_remaining+' | Age: '+r.session_age+'s','ok');else showResult('account-result','No active session','warn');}
function showResult(id,msg,type){const el=document.getElementById(id);el.style.display='block';el.innerHTML='<span class="'+type+'">'+msg+'</span>';}
async function loadHistory(){
const r=await api('/history');const el=document.getElementById('history-table');
if(!r.queries||!r.queries.length){el.innerHTML='<p style="color:#64748b">No queries yet</p>';return}
let html='<table><thead><tr><th>Time</th><th>Thread</th><th>Query</th><th>Mode</th><th>Model</th><th>Result</th></tr></thead><tbody>';
r.queries.slice(-50).reverse().forEach(q=>{
const t=q.timestamp?new Date(q.timestamp).toLocaleString():'';
const cls=q.success?'ok':'err';
const tid=q.thread_id?q.thread_id.substring(0,8)+'...':'--';
html+='<tr><td class="mono">'+t+'</td><td class="mono">'+tid+'</td><td>'+esc(q.query)+'</td><td>'+q.mode+'</td><td>'+q.model+'</td><td class="'+cls+'">'+(q.success?q.answer_len+' chars':'failed')+'</td></tr>';
});html+='</tbody></table>';el.innerHTML=html;
}
function esc(s){return s?s.replace(/&/g,'&amp;').replace(/</g,'&lt;').replace(/>/g,'&gt;'):''}
async function loadModels(){
const r=await api('/models');const mt=document.getElementById('models-table');const lt=document.getElementById('labs-table');
let html='';
if(r.search_models)Object.entries(r.search_models).forEach(([mode,models])=>{
models.forEach((m,i)=>{html+='<tr><td>'+(i===0?mode:'')+'</td><td>'+(m.name||'default')+'</td><td class="mono">'+m.internal+'</td><td>'+(m.auth?'Yes':'No')+'</td></tr>';});
});mt.innerHTML=html;
let lhtml='';if(r.labs_models)r.labs_models.forEach(m=>{lhtml+='<tr><td>'+m+'</td><td class="mono">WebSocket</td></tr>';});lt.innerHTML=lhtml;
}
async function loadThreads(){
const r=await api('/threads');const el=document.getElementById('threads-table');
if(!r.threads||!Object.keys(r.threads).length){el.innerHTML='<p style="color:#64748b">No threads yet</p>';return}
let html='<table><thead><tr><th>ID</th><th>Mode</th><th>Backend UUID</th><th>Messages</th><th>Created</th></tr></thead><tbody>';
Object.entries(r.threads).forEach(([id,t])=>{
const tid=id.substring(0,12)+'...';
const buid=t.backend_uuid?t.backend_uuid.substring(0,12)+'...':'<span class="warn">expired</span>';
const msgCount=Array.isArray(t.messages)?t.messages.length:'?';
html+='<tr><td class="mono">'+tid+'</td><td>'+t.mode+'</td><td class="mono">'+buid+'</td><td>'+msgCount+'</td><td class="mono">'+t.created_at+'</td></tr>';
});html+='</tbody></table>';el.innerHTML=html;
}
async function loadLogs(){
const r=await api('/logs');const el=document.getElementById('log-output');
if(r.logs&&r.logs.length)el.innerHTML=r.logs.map(l=>'<span class="'+(l.level||'info').toLowerCase()+'">['+l.level+']</span> '+esc(l.message)).join('\\n');
else el.innerHTML='No logs available';
}
refreshStatus();loadHistory();loadModels();
setInterval(refreshStatus,30000);
</script>
</body>
</html>"""


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

app = FastAPI(title="PPLX v2 - Perplexity AI Proxy", version="2.0.0")


class AuthMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        path = request.url.path

        if path.startswith("/v1/") or path == "/v1/models":
            if API_KEY:
                auth = request.headers.get("Authorization", "")
                key = auth.replace("Bearer ", "") if auth.startswith("Bearer ") else request.query_params.get("key", "")
                if key != API_KEY:
                    return JSONResponse({"error": {"message": "Invalid API key", "type": "auth_error"}}, status_code=401)

        if path == "/admin/login" and ADMIN_PASSWORD:
            pwd = request.query_params.get("password", "")
            if pwd == ADMIN_PASSWORD:
                response = HTMLResponse('<html><body><script>localStorage.setItem("admin_token","' + pwd + '");window.location="/"</script></body></html>')
                response.set_cookie("admin_token", pwd, max_age=86400, httponly=False, path="/")
                return response
            return HTMLResponse("<h3>Wrong password</h3>", status_code=401)

        if (path.startswith("/admin") or path == "/") and path != "/admin/login":
            if ADMIN_PASSWORD:
                auth = request.headers.get("Authorization", "")
                pwd = auth.replace("Bearer ", "") if auth.startswith("Bearer ") else request.query_params.get("password", "") or request.cookies.get("admin_token", "")
                if pwd != ADMIN_PASSWORD:
                    if path == "/":
                        return HTMLResponse('<html><body><h3>Admin Login</h3><form method="GET" action="/admin/login"><input name="password" type="password" placeholder="Password"><button>Login</button></form></body></html>', status_code=401)
                    return JSONResponse({"error": "Unauthorized"}, status_code=401)

        return await call_next(request)


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


# ─── Endpoints ─────────────────────────────────────────────────────────────────

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


@app.get("/admin/api/status")
async def api_status():
    stored = load_cookies()
    if stored and stored.get("cookies"):
        age = int(time.time() - stored.get("saved_at", 0))
        return {
            "cookies_active": True,
            "copilot_remaining": stored.get("copilot_remaining", 0),
            "file_upload_remaining": stored.get("file_upload_remaining", 0),
            "session_age": age,
        }
    return {"cookies_active": False, "copilot_remaining": -1, "file_upload_remaining": -1, "session_age": 0}


@app.post("/admin/api/account/create")
async def api_create_account():
    try:
        future = EXECUTOR.submit(auto_create_account)
        cookies = future.result(timeout=ACCOUNT_CREATE_TIMEOUT)
        client = Client(cookies)
        save_cookies(cookies, client.copilot, client.file_upload)
        return {"ok": True, "copilot_remaining": client.copilot, "file_upload_remaining": client.file_upload}
    except Exception as e:
        logger.error(f"Account creation failed: {e}")
        return {"ok": False, "error": str(e)}


@app.get("/admin/api/account/check")
async def api_check_account():
    stored = load_cookies()
    if stored and stored.get("cookies"):
        age = int(time.time() - stored.get("saved_at", 0))
        return {
            "ok": True,
            "copilot_remaining": stored.get("copilot_remaining", 0),
            "file_upload_remaining": stored.get("file_upload_remaining", 0),
            "session_age": age,
        }
    return {"ok": False, "error": "No active session"}


@app.post("/admin/api/session/clear")
async def api_clear_session():
    try:
        if COOKIES_FILE.exists():
            COOKIES_FILE.unlink()
        return {"ok": True}
    except Exception as e:
        return {"ok": False, "error": str(e)}


@app.get("/admin/api/history")
async def api_history():
    try:
        if HISTORY_FILE.exists():
            with open(HISTORY_FILE, "r") as f:
                return {"queries": json.load(f)}
        return {"queries": []}
    except Exception:
        return {"queries": []}


@app.get("/admin/api/threads")
async def api_threads():
    return {"threads": _load_threads()}


@app.delete("/admin/api/threads/{thread_id}")
async def api_delete_thread(thread_id: str):
    threads = _load_threads()
    if thread_id in threads:
        del threads[thread_id]
        _save_threads(threads)
        return {"ok": True}
    return {"ok": False, "error": "Thread not found"}


@app.get("/admin/api/models")
async def api_models():
    search_models = {}
    for mode, models in MODEL_MAPPINGS.items():
        search_models[mode] = [
            {"name": k if k else "default", "internal": v, "auth": mode != "auto"}
            for k, v in models.items()
        ]
    return {"search_models": search_models, "labs_models": LABS_MODELS}


@app.get("/admin/api/logs")
async def api_logs():
    try:
        if not LOG_FILE.exists():
            return {"logs": []}
        lines = []
        with open(LOG_FILE, "r") as f:
            for line in f.readlines()[-200:]:
                line = line.strip()
                if not line:
                    continue
                try:
                    parts = line.split(" - ", 3)
                    if len(parts) >= 4:
                        lines.append({"timestamp": parts[0], "logger": parts[1], "level": parts[2].strip(), "message": parts[3].strip()})
                    else:
                        lines.append({"level": "INFO", "message": line})
                except Exception:
                    lines.append({"level": "INFO", "message": line})
        return {"logs": list(reversed(lines))}
    except Exception:
        return {"logs": []}


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

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

    sources = _extract_sources(messages)
    mode = MODE_MAP.get(model, "auto")
    model_val = model if model in VALID_MODEL_NAMES else None

    if not thread_id and len(messages) > 2:
        thread_id = str(uuid.uuid4())
        context, contextual_query = _build_context_query(messages)
        if context:
            query = contextual_query

    if thread_id and not _get_thread(thread_id):
        _update_thread(thread_id, {"mode": mode})

    req_id = "chatcmpl-{}".format(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, sources, 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 "data: {}\n\n".format(_build_openai_chunk(req_id, created, model, {"content": delta_text}))
                yield "data: {}\n\n".format(_build_openai_chunk(req_id, created, model, {}, "stop"))
                yield "data: [DONE]\n\n"
                if thread_id:
                    _update_thread(thread_id, {"messages": messages + [{"role": "assistant", "content": full_text}]})
                _log_query(query, mode, model_val, len(full_text), True, thread_id)
            except Exception as e:
                error_chunk = _build_openai_chunk(req_id, created, model, {"content": "Error: {}".format(str(e))}, "stop")
                yield "data: {}\n\n".format(error_chunk)
                yield "data: [DONE]\n\n"
                _log_query(query, mode, model_val, 0, False, thread_id)
        return StreamingResponse(generate(), media_type="text/event-stream")
    else:
        try:
            result = search_sync(query, mode, model_val, sources, language, "", thread_id=thread_id)
            if thread_id:
                _update_thread(thread_id, {"messages": messages + [{"role": "assistant", "content": result}]})
            _log_query(query, mode, model_val, len(result), True, thread_id)
            response = _build_openai_response(req_id, created, model, result)
            if thread_id:
                response["thread_id"] = thread_id
            return JSONResponse(response)
        except PerplexityError as e:
            _log_query(query, mode, model_val, 0, False, thread_id)
            return JSONResponse({"error": {"message": str(e), "type": type(e).__name__}}, status_code=500)
        except Exception as e:
            _log_query(query, mode, model_val, 0, False, thread_id)
            return JSONResponse({"error": {"message": str(e), "type": "server_error"}}, status_code=500)


@app.post("/v1/responses")
async def openai_responses(request: Request):
    body = await request.json()
    input_data = body.get("input", [])
    model = body.get("model", "auto")
    stream = body.get("stream", False)
    language = body.get("language", "en-US")
    thread_id = body.get("thread_id")

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

    system_text = _extract_system_from_input(input_data)
    if "scholar" in system_text.lower() or "academic" in system_text.lower():
        sources = "scholar"
    elif "news" in system_text.lower():
        sources = "news"
    else:
        sources = "web"

    mode = MODE_MAP.get(model, "auto")
    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})

    resp_id = "resp_{}".format(uuid.uuid4().hex[:24])
    item_id = "item_{}".format(uuid.uuid4().hex[:24])
    output_id = "item_{}".format(uuid.uuid4().hex[:24])
    created = int(time.time())

    if stream:
        async def generate():
            full_text = ""
            yield _build_responses_sse_event("response.created", {
                "type": "response.created",
                "response": {"id": resp_id, "object": "response", "created_at": created, "model": model, "status": "in_progress", "output": []}
            })
            yield _build_responses_sse_event("response.output_item.added", {
                "type": "response.output_item.added", "output_index": 0,
                "item": {"id": item_id, "type": "message", "role": "assistant", "status": "in_progress", "content": []}
            })
            yield _build_responses_sse_event("response.content_part.added", {
                "type": "response.content_part.added", "output_index": 0, "content_index": 0,
                "part": {"type": "output_text", "text": "", "annotations": []}
            })
            try:
                for chunk_text in stream_search(query, mode, model_val, sources, 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 _build_responses_sse_event("response.output_text.delta", {
                            "type": "response.output_text.delta", "output_index": 0, "content_index": 0, "delta": delta_text
                        })
                yield _build_responses_sse_event("response.output_text.done", {
                    "type": "response.output_text.done", "output_index": 0, "content_index": 0, "text": full_text
                })
                yield _build_responses_sse_event("response.content_part.done", {
                    "type": "response.content_part.done", "output_index": 0, "content_index": 0,
                    "part": {"type": "output_text", "text": full_text, "annotations": []}
                })
                yield _build_responses_sse_event("response.output_item.done", {
                    "type": "response.output_item.done", "output_index": 0,
                    "item": {"id": item_id, "type": "message", "role": "assistant", "status": "completed", "content": [{"type": "output_text", "text": full_text, "annotations": []}]}
                })
                yield _build_responses_sse_event("response.completed", {
                    "type": "response.completed",
                    "response": {
                        "id": resp_id, "object": "response", "created_at": created, "model": model,
                        "status": "completed",
                        "output": [{"id": item_id, "type": "message", "role": "assistant", "status": "completed", "content": [{"type": "output_text", "text": full_text, "annotations": []}]}],
                        "usage": {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0}
                    }
                })
                if thread_id:
                    _update_thread(thread_id, {"messages": [{"role": "user", "content": query}, {"role": "assistant", "content": full_text}]})
                _log_query(query, mode, model_val, len(full_text), True, thread_id)
            except Exception as e:
                yield _build_responses_sse_event("error", {"type": "error", "message": str(e)})
                _log_query(query, mode, model_val, 0, False, thread_id)

        return StreamingResponse(generate(), media_type="text/event-stream")
    else:
        try:
            result = search_sync(query, mode, model_val, sources, language, "", thread_id=thread_id)
            if thread_id:
                _update_thread(thread_id, {"messages": [{"role": "user", "content": query}, {"role": "assistant", "content": result}]})
            _log_query(query, mode, model_val, len(result), True, thread_id)
            return JSONResponse({
                "id": resp_id, "object": "response", "created_at": created, "model": model,
                "status": "completed",
                "output": [{"id": output_id, "type": "message", "role": "assistant", "status": "completed", "content": [{"type": "output_text", "text": result, "annotations": []}]}],
                "usage": {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0}
            })
        except PerplexityError as e:
            _log_query(query, mode, model_val, 0, False, thread_id)
            return JSONResponse({"error": {"message": str(e), "type": type(e).__name__}}, status_code=500)
        except Exception as e:
            _log_query(query, mode, model_val, 0, False, thread_id)
            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, mode_models in MODEL_MAPPINGS.items():
        for k in mode_models:
            if k is not None:
                models.append({"id": k, "object": "model", "owned_by": "perplexity"})
    for lab_model in LABS_MODELS:
        models.append({"id": lab_model, "object": "model", "owned_by": "perplexity-labs"})
    return JSONResponse({"object": "list", "data": models})


@app.get("/health")
async def health():
    stored = load_cookies()
    copilot = stored.get("copilot_remaining", -1) if stored else -1
    threads = _load_threads()
    return {
        "status": "ok",
        "version": "2.0.0",
        "cookies_active": bool(stored and stored.get("cookies")),
        "copilot_remaining": copilot,
        "active_threads": len(threads),
    }


@app.post("/mcp")
async def mcp_endpoint(request: Request):
    body = await request.json()
    tool = body.get("tool", "ask")
    query = body.get("query", body.get("input", ""))
    thread_id = body.get("thread_id")
    if not query:
        return JSONResponse({"error": "No query provided"}, status_code=400)

    mode_map = {
        "ask": "auto",
        "search": "pro",
        "reason": "reasoning",
        "research": "deep research",
    }
    mode = mode_map.get(tool, "auto")

    if not thread_id:
        thread_id = str(uuid.uuid4())

    try:
        result = search_sync(query, mode, thread_id=thread_id)
        return {"tool": tool, "mode": mode, "result": result, "thread_id": thread_id}
    except Exception as e:
        return JSONResponse({"error": str(e)}, status_code=500)


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