from flask import Flask, request, jsonify, Response, stream_with_context
import requests
import json
import sys

# Force UTF-8 output
try:
    sys.stdout.reconfigure(encoding="utf-8")
except Exception:
    pass

app = Flask(__name__)

API_URL = "https://ml-gateway.voiceflow.com/v1alpha1/generation/global-instruction/stream"

HEADERS = {
    "Authorization": "Bearer eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJ1c2VySUQiOjMwMDU1OCwiZW1haWwiOiJhbGhvb3JzaG9wcEBnbWFpbC5jb20iLCJuYW1lIjoib21hciBudXdhcmEiLCJpbnRlcm5hbEFkbWluIjpmYWxzZSwiaXNzIjoiaHR0cHM6Ly9hdXRoLWFwaS52b2ljZWZsb3cuY29tIiwiaWF0IjoxNzgwOTUzMzA2fQ.lrVXj60S5tGG7Ld1R_pCZ3dRopVsB-V7pZlORBp7l78",
    "Content-Type": "application/json",
    "Accept": "text/event-stream",
}

PROJECT_ID = "6a27315d4b4451fda7e0d81f"
VERSION_ID = "6a27315d4b4451fda7e0d820"
WORKSPACE_ID = "95kwg6mRkO"

START_TAG = "<AI>"
END_TAG = "</AI>"


def repair_text(text):
    """Attempts to fix common UTF-8/Latin1 corruption."""
    if not text:
        return text
    try:
        repaired = text.encode("latin1").decode("utf-8")
        if repaired.count("") <= text.count(""):
            text = repaired
    except Exception:
        pass
    return text


def build_prompt(messages):
    conversation = []
    for msg in messages:
        role = msg.get("role", "").lower()
        content = msg.get("content", "")
        if not content:
            continue
        content = repair_text(content)
        if role == "system":
            conversation.append(f"SYSTEM: {content}")
        elif role == "user":
            conversation.append(f"USER: {content}")
        elif role == "assistant":
            conversation.append(f"AI: {content}")
    conversation.append("ASSISTANT:")
    return "\n\n".join(conversation)


def create_payload(messages):
    return {
        "locales": ["en-US"],
        "variables": [],
        "projectID": PROJECT_ID,
        "versionID": VERSION_ID,
        "workspaceID": WORKSPACE_ID,
        "userInstruction": "okay you have two tasks create instructions and response to user but as XML :<AI>youe response to user here with all XMLs from the system</AI> THEN <Instruction>your Instruction here</Instruction>\n\nCHAT :" + build_prompt(messages),
        "previousUserInstruction": ""
    }


@app.route("/chat", methods=["POST"])
def chat():
    data = request.get_json(silent=True)

    if not data:
        return jsonify({"error": "Invalid JSON"}), 400

    print("\n" + "=" * 60)
    print("📥 INCOMING REQUEST:")
    print(json.dumps(data, indent=2, ensure_ascii=False))
    print("=" * 60 + "\n")

    messages = data.get("messages")
    if not messages:
        return jsonify({"error": "messages required"}), 400

    stream = data.get("stream", False)
    payload = create_payload(messages)

    if stream:
        @stream_with_context
        def generate():
            buffer = ""
            inside_ai = False
            ai_started = False  # Tracks if we've seen <AI> or started fallback streaming
            try:
                with requests.post(
                    API_URL, headers=HEADERS, json=payload, stream=True, timeout=120
                ) as response:
                    response.raise_for_status()
                    response.encoding = "utf-8"

                    print("📤 STREAM OUTPUT (Filtered):")
                    for line in response.iter_lines(decode_unicode=True):
                        if not line or not line.startswith("data: "):
                            continue
                        try:
                            event = json.loads(line[6:])
                            chunk = event.get("output", "")
                            if not chunk:
                                continue

                            buffer += chunk

                            # --- START TAG DETECTION ---
                            if not inside_ai and START_TAG in buffer:
                                idx = buffer.index(START_TAG) + len(START_TAG)
                                buffer = buffer[idx:]
                                inside_ai = True
                                ai_started = True

                            # --- FALLBACK: If model skips <AI> tag entirely ---
                            if not ai_started and len(buffer) > 50:
                                inside_ai = True
                                ai_started = True

                            # --- INSIDE AI CONTENT ---
                            if inside_ai:
                                if END_TAG in buffer:
                                    # ✅ STOP STREAMING IMMEDIATELY after </AI>
                                    idx = buffer.index(END_TAG)
                                    final_chunk = buffer[:idx]
                                    if final_chunk:
                                        print(final_chunk, end="", flush=True)
                                        yield (
                                            "data: "
                                            + json.dumps({"content": final_chunk}, ensure_ascii=False)
                                            + "\n\n"
                                        )
                                    print("\n✅ </AI> REACHED - STOPPING STREAM\n")
                                    yield "data: [DONE]\n\n"
                                    return  # ⛔ Hard stop: ignore all remaining chunks

                                else:
                                    safe_len = max(0, len(buffer) - len(END_TAG))
                                    if safe_len > 0:
                                        to_yield = buffer[:safe_len]
                                        buffer = buffer[safe_len:]
                                        print(to_yield, end="", flush=True)
                                        yield (
                                            "data: "
                                            + json.dumps({"content": to_yield}, ensure_ascii=False)
                                            + "\n\n"
                                        )

                        except Exception as e:
                            print(f"\n⚠️ Chunk Error: {e}")

                    print("\n✅ STREAM COMPLETE\n")
                    yield "data: [DONE]\n\n"

            except Exception as e:
                error_msg = str(e)
                print(f"\n❌ STREAM ERROR: {error_msg}\n")
                yield "data: " + json.dumps({"error": error_msg}) + "\n\n"

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

    # Non-streaming path
    try:
        full_response = ""
        buffer = ""
        inside_ai = False
        ai_started = False
        print("📤 NON-STREAM OUTPUT (Filtered):")

        with requests.post(
            API_URL, headers=HEADERS, json=payload, stream=True, timeout=120
        ) as response:
            response.raise_for_status()
            response.encoding = "utf-8"

            for line in response.iter_lines(decode_unicode=True):
                if not line or not line.startswith("data: "):
                    continue
                try:
                    event = json.loads(line[6:])
                    chunk = event.get("output", "")
                    if not chunk:
                        continue

                    buffer += chunk

                    if not inside_ai and START_TAG in buffer:
                        idx = buffer.index(START_TAG) + len(START_TAG)
                        buffer = buffer[idx:]
                        inside_ai = True
                        ai_started = True

                    if not ai_started and len(buffer) > 50:
                        inside_ai = True
                        ai_started = True

                    if inside_ai:
                        if END_TAG in buffer:
                            idx = buffer.index(END_TAG)
                            full_response += buffer[:idx]
                            break  # ⛔ Stop reading after </AI>
                        else:
                            safe_len = max(0, len(buffer) - len(END_TAG))
                            if safe_len > 0:
                                full_response += buffer[:safe_len]
                                buffer = buffer[safe_len:]

                except Exception:
                    pass

        full_response = repair_text(full_response)
        print(full_response)
        print("\n✅ RESPONSE COMPLETE\n")

        return jsonify({"role": "assistant", "content": full_response})

    except requests.exceptions.RequestException as e:
        print(f"\n❌ REQUEST ERROR: {e}\n")
        return jsonify({"error": str(e)}), 500
    except Exception as e:
        print(f"\n❌ ERROR: {e}\n")
        return jsonify({"error": str(e)}), 500


@app.route("/", methods=["GET"])
def home():
    return jsonify({"status": "online", "endpoint": "/chat"})


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860, threaded=True, debug=True)