"""
MindStudio WebSocket Streaming Client — Flask API Endpoint
FIXED: Race condition — AI was answering previous message (buffer drain fix)
FIXED: Per-thread continuation tokens
FIXED: WebSocket compatibility
"""

import asyncio
import json
import httpx
import websockets
import os
import threading

from flask import Flask, request, Response, jsonify, stream_with_context

app = Flask(__name__)

# ─────────────────────────────────────────────
# CONFIGURATION
# ─────────────────────────────────────────────

ORG_ID = os.getenv("MINDSTUDIO_ORG_ID", "9235df37-94ad-47a2-92f7-bbf28ea8d4d1")
APP_ID = os.getenv("MINDSTUDIO_APP_ID", "75a42c47-67b3-4c9d-a23f-65998823d266")
INITIAL_CONTINUATION_TOKEN = os.getenv(
    "MINDSTUDIO_CONTINUATION_TOKEN",
    "155b2f71-9738-4260-8849-46e3eb0dd384::2b71a3fd-f4f0-4439-8163-60cd53d28218"
)
TARGET_MODEL = "claude-4-7-opus"
WS_URL = "wss://api-socket.mindstudio.ai/"

# ─────────────────────────────────────────────
# GLOBAL STATE
# ─────────────────────────────────────────────

_clients: dict = {}
_clients_lock = threading.Lock()


# ─────────────────────────────────────────────
# HELPERS
# ─────────────────────────────────────────────

def safe_get_model(debug_info: dict):
    model_settings = debug_info.get("modelSettings")
    if isinstance(model_settings, dict):
        return model_settings.get("model")
    return None


def check_ws_open(ws) -> bool:
    if ws is None:
        return False
    try:
        if hasattr(ws, 'close_code'):
            return ws.close_code is None
        if hasattr(ws, 'closed'):
            return not ws.closed
        if hasattr(ws, 'state'):
            import websockets.connection
            return ws.state == websockets.connection.State.OPEN
        return True
    except Exception:
        return False


# ─────────────────────────────────────────────
# MINDSTUDIO CLIENT
# ─────────────────────────────────────────────

class MindStudioClient:
    def __init__(self, access_token: str, model: str = TARGET_MODEL):
        self.access_token = access_token
        self.model = model

        # Per-thread continuation tokens
        self._continuation_tokens: dict = {}
        self._tokens_lock = threading.Lock()

        self.ws = None
        self.loop = asyncio.new_event_loop()

        self._thread = threading.Thread(target=self._run_loop, daemon=True)
        self._thread.start()

    def _run_loop(self):
        asyncio.set_event_loop(self.loop)
        self.loop.run_forever()

    def connect_sync(self):
        future = asyncio.run_coroutine_threadsafe(self._connect(), self.loop)
        future.result(timeout=20)

    async def _connect(self):
        print(f"[WS] Connecting to {WS_URL} ...")
        extra_headers = {
            "Origin": "https://app.mindstudio.ai",
            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64)",
        }
        self.ws = await websockets.connect(
            WS_URL,
            subprotocols=["auth", self.access_token],
            additional_headers=extra_headers,
            ping_interval=None,
            ping_timeout=None,
        )
        print(f"[WS] Connected! Type: {type(self.ws).__name__}")
        print(f"[WS] Subprotocol: {self.ws.subprotocol}")

        has_closed    = hasattr(self.ws, 'closed')
        has_closecode = hasattr(self.ws, 'close_code')
        has_state     = hasattr(self.ws, 'state')
        print(f"[WS] Attrs → closed={has_closed} close_code={has_closecode} state={has_state}")

    def is_connected(self) -> bool:
        return check_ws_open(self.ws)

    def get_continuation_token(self, thread_id: str) -> str:
        with self._tokens_lock:
            return self._continuation_tokens.get(thread_id, INITIAL_CONTINUATION_TOKEN)

    def set_continuation_token(self, thread_id: str, token: str):
        with self._tokens_lock:
            self._continuation_tokens[thread_id] = token
            print(f"[TOKEN] Thread {thread_id[:8]}... token updated")

    # ─────────────────────────────────────────
    # ✅ FIX: Drain leftover WS buffer
    # ─────────────────────────────────────────
    async def _drain_ws_buffer(self):
        """Flush any leftover WebSocket messages from a previous response."""
        flushed = 0
        try:
            while True:
                raw = await asyncio.wait_for(self.ws.recv(), timeout=0.05)
                flushed += 1
                print(f"[DRAIN] Flushed leftover msg #{flushed}: {str(raw)[:80]}...")
        except asyncio.TimeoutError:
            pass
        except Exception:
            pass
        if flushed:
            print(f"[DRAIN] Total flushed: {flushed} leftover message(s)")

    # ─────────────────────────────────────────
    # ✅ FIX: Post FIRST, then listen
    # ─────────────────────────────────────────
    def send_message_stream(self, message: str, thread_id: str):
        import queue

        chunk_queue = queue.Queue()
        done_event  = threading.Event()

        async def _run():
            try:
                # ✅ STEP 1: Drain any leftover buffer from previous message
                await self._drain_ws_buffer()

                # ✅ STEP 2: Post the new message
                await self._post_message(message, thread_id)

                # ✅ STEP 3: Small wait for server to process
                await asyncio.sleep(0.15)

                # ✅ STEP 4: Listen for THIS message's response only
                await self._listen_for_response(chunk_queue, done_event, thread_id)

            except Exception as e:
                print(f"[ERROR] send_message_stream: {e}")
                chunk_queue.put(f"[ERROR] {e}")
                done_event.set()

        asyncio.run_coroutine_threadsafe(_run(), self.loop)

        while not done_event.is_set() or not chunk_queue.empty():
            try:
                chunk = chunk_queue.get(timeout=0.5)
                yield f"data: {json.dumps({'text': chunk})}\n\n"
            except Exception:
                continue

        yield f"data: {json.dumps({'done': True})}\n\n"

    async def _post_message(self, message: str, thread_id: str):
        post_url = f"https://v1.mindstudio-api.com/v1/apps/load/{APP_ID}/threads/post"

        continuation_token = self.get_continuation_token(thread_id)
        print(f"[POST] Using token: {continuation_token[:30]}... for thread: {thread_id[:8]}...")

        payload = {
            "threadId": thread_id,
            "continuationToken": continuation_token,
            "payload": {
                "message": message,
                "imageUrl": "",
                "modelCandidates": [self.model],
            },
        }
        headers = {
            "Authorization": f"Bearer: {self.access_token}",
            "X-Organization-Id": ORG_ID,
            "Content-Type": "application/json",
            "Accept": "application/json",
            "Origin": "https://app.mindstudio.ai",
            "User-Agent": "Mozilla/5.0",
            "Referer": "https://app.mindstudio.ai/",
        }
        async with httpx.AsyncClient() as client:
            resp = await client.post(post_url, json=payload, headers=headers)
        print(f"[POST] Status: {resp.status_code}")
        if resp.status_code not in (200, 204):
            print(f"[POST] Body: {resp.text}")

    async def _listen_for_response(self, chunk_queue, done_event, thread_id: str):
        full_response = ""
        previous_content_length = 0
        system_post_id = None
        detected_model = None

        while True:
            try:
                raw_message = await asyncio.wait_for(self.ws.recv(), timeout=30.0)
            except asyncio.TimeoutError:
                continue
            except websockets.exceptions.ConnectionClosed as e:
                print(f"[WS] Connection closed: {e}")
                break
            except Exception as e:
                print(f"[WS] Recv error: {e}")
                break

            try:
                event = json.loads(raw_message)
            except json.JSONDecodeError:
                continue

            event_type = event.get("type", "")

            if event_type == "Apps/ThreadContinuationActionUpdated":
                data = event.get("appThreadContinuationActionUpdated", {})
                new_token = data.get("continuationAction", {}).get("token")
                if new_token:
                    self.set_continuation_token(thread_id, new_token)

                underlying = (
                    data.get("continuationAction", {})
                    .get("step", {})
                    .get("underlyingModel", {})
                )
                if isinstance(underlying, dict) and underlying.get("model"):
                    if not detected_model:
                        detected_model = underlying["model"]

            elif event_type == "Apps/ThreadPostsCreated":
                posts = event.get("appThreadPostsCreated", {}).get("posts", [])
                for post in posts:
                    chat_msg = post.get("chatMessage", {})
                    if chat_msg.get("source") == "system":
                        system_post_id = post.get("id")

                        debug_info = chat_msg.get("_debugInfo", {})
                        if isinstance(debug_info, dict):
                            model = safe_get_model(debug_info)
                            if model:
                                detected_model = model

                        content = chat_msg.get("content", "")
                        is_in_progress = chat_msg.get("isInProgress", True)

                        if content and not is_in_progress:
                            chunk_queue.put(content)
                            done_event.set()
                            return
                        elif content:
                            chunk_queue.put(content)
                            previous_content_length = len(content)
                            full_response = content

            elif event_type == "Apps/ThreadPostUpdated":
                post_data = event.get("appThreadPostUpdated", {})
                new_post  = post_data.get("newPost", {})
                post_id   = post_data.get("postId")
                chat_msg  = new_post.get("chatMessage", {})

                if chat_msg.get("source") != "system":
                    continue
                if system_post_id and post_id != system_post_id:
                    continue

                debug_info = chat_msg.get("_debugInfo", {})
                if isinstance(debug_info, dict):
                    model = safe_get_model(debug_info)
                    if model and not detected_model:
                        detected_model = model

                current_content = chat_msg.get("content", "")
                is_in_progress  = chat_msg.get("isInProgress", True)

                if len(current_content) > previous_content_length:
                    delta = current_content[previous_content_length:]
                    chunk_queue.put(delta)
                    previous_content_length = len(current_content)
                    full_response = current_content

                if not is_in_progress:
                    print(f"[DONE] Model: {detected_model or 'unknown'}")
                    done_event.set()
                    return

        done_event.set()

    async def _close(self):
        if self.ws:
            await self.ws.close()


# ─────────────────────────────────────────────
# HELPER: get or create client
# ─────────────────────────────────────────────

def get_or_create_client(access_token: str) -> MindStudioClient:
    with _clients_lock:
        client = _clients.get(access_token)
        if client and client.is_connected():
            print("[CLIENT] Reusing existing connected client")
            return client
        print("[CLIENT] Creating new client and connecting...")
        client = MindStudioClient(access_token=access_token)
        client.connect_sync()
        _clients[access_token] = client
        return client


# ─────────────────────────────────────────────
# ROUTES
# ─────────────────────────────────────────────

@app.route("/health", methods=["GET"])
def health():
    return jsonify({
        "status": "ok",
        "model": TARGET_MODEL,
        "websockets_version": websockets.__version__,
    }), 200


@app.route("/connect", methods=["GET"])
def connect():
    auth = request.headers.get("Authorization", "")
    if not auth.startswith("Bearer "):
        return jsonify({"error": "Missing or invalid Authorization header"}), 401

    access_token = auth.removeprefix("Bearer ").strip()

    try:
        client = get_or_create_client(access_token)
        ws_type = type(client.ws).__name__ if client.ws else "None"
        return jsonify({
            "status": "connected",
            "model": TARGET_MODEL,
            "ws_object_type": ws_type,
        }), 200
    except Exception as e:
        print(f"[CONNECT ERROR] {type(e).__name__}: {e}")
        return jsonify({"error": str(e)}), 500


@app.route("/chat", methods=["POST"])
def chat():
    auth = request.headers.get("Authorization", "")
    if not auth.startswith("Bearer "):
        return jsonify({"error": "Missing or invalid Authorization header"}), 401

    access_token = auth.removeprefix("Bearer ").strip()

    data = request.get_json(silent=True)
    if not data:
        return jsonify({"error": "Invalid JSON body"}), 400

    message   = data.get("message", "").strip()
    thread_id = data.get("threadId", "").strip()

    if not message:
        return jsonify({"error": "Missing 'message' field"}), 400
    if not thread_id:
        return jsonify({"error": "Missing 'threadId' field"}), 400

    try:
        client = get_or_create_client(access_token)
    except Exception as e:
        return jsonify({"error": f"WebSocket connection failed: {e}"}), 500

    def generate():
        yield f"data: {json.dumps({'status': 'connected', 'model': TARGET_MODEL})}\n\n"
        for chunk in client.send_message_stream(message, thread_id):
            yield chunk

    return Response(
        stream_with_context(generate()),
        mimetype="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "X-Accel-Buffering": "no",
            "Connection": "keep-alive",
        },
    )


# ─────────────────────────────────────────────
# ENTRY POINT
# ─────────────────────────────────────────────

if __name__ == "__main__":
    print("="*60)
    print(f"  MindStudio Flask API  |  Model: {TARGET_MODEL}")
    print(f"  websockets version: {websockets.__version__}")
    print("="*60)
    app.run(host="0.0.0.0", port=7860, debug=False, threaded=True)