#!/usr/bin/env python3
"""
mock_chatgpt.py - ChatGPT-like interface with multi-model routing
Uses FlowtTorch backend with intelligent model selection
"""
import os, json, hashlib, time
from typing import List, Dict, Any
from fastapi import FastAPI, HTTPException
from fastapi.responses import HTMLResponse, StreamingResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
import httpx

app = FastAPI(title="MockChatGPT", version="0.1.0")

# FlowtTorch backend
FT_BASE = os.environ.get("FT", "http://127.0.0.1:8089")

# Model routing map
MODEL_ROUTES = {
    "code": "flowtorch/code",       # Code generation/analysis
    "chat": "flowtorch/chat",       # General conversation
    "general": "flowtorch/braid",   # Default catch-all
}

# Simple keyword-based routing
CODE_KEYWORDS = {"function", "class", "import", "def", "return", "loop", "array", "variable"}
CHAT_KEYWORDS = {"hello", "hi", "how", "what", "why", "when", "where", "who", "explain"}


class Message(BaseModel):
    role: str
    content: str


class ChatRequest(BaseModel):
    messages: List[Message]
    max_tokens: int = 256
    helix: Dict[str, Any] = {}


def tag_prompt(content: str) -> str:
    """Simple keyword-based tagging for model selection."""
    content_lower = content.lower()

    # Check for code keywords
    code_score = sum(1 for kw in CODE_KEYWORDS if kw in content_lower)
    chat_score = sum(1 for kw in CHAT_KEYWORDS if kw in content_lower)

    if code_score > chat_score and code_score > 0:
        return "code"
    elif chat_score > 0:
        return "chat"
    else:
        return "general"


@app.get("/", response_class=HTMLResponse)
async def ui():
    """Simple chat UI."""
    return """
<!DOCTYPE html>
<html>
<head>
    <title>MockChatGPT</title>
    <style>
        body {
            font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, sans-serif;
            max-width: 800px;
            margin: 40px auto;
            padding: 20px;
            background: #f7f7f8;
        }
        #chat {
            background: white;
            border-radius: 8px;
            padding: 20px;
            min-height: 400px;
            max-height: 600px;
            overflow-y: auto;
            margin-bottom: 20px;
            box-shadow: 0 2px 8px rgba(0,0,0,0.1);
        }
        .msg {
            margin: 10px 0;
            padding: 12px;
            border-radius: 6px;
        }
        .user {
            background: #e3f2fd;
            text-align: right;
        }
        .assistant {
            background: #f5f5f5;
        }
        .meta {
            font-size: 0.8em;
            color: #666;
            margin-top: 5px;
        }
        #input-box {
            display: flex;
            gap: 10px;
        }
        #prompt {
            flex: 1;
            padding: 12px;
            border: 1px solid #ddd;
            border-radius: 6px;
            font-size: 14px;
        }
        #send {
            padding: 12px 24px;
            background: #10a37f;
            color: white;
            border: none;
            border-radius: 6px;
            cursor: pointer;
            font-weight: 600;
        }
        #send:hover { background: #0d8c6d; }
        #send:disabled {
            background: #ccc;
            cursor: not-allowed;
        }
        .receipt {
            background: #fff9e6;
            border-left: 3px solid #ffa500;
            padding: 8px;
            margin-top: 8px;
            font-size: 0.75em;
            font-family: monospace;
        }
    </style>
</head>
<body>
    <h1>🧬 MockChatGPT</h1>
    <p>Multi-model routing with SeedOps receipts</p>
    <div id="chat"></div>
    <div id="input-box">
        <input id="prompt" type="text" placeholder="Type a message..." />
        <button id="send">Send</button>
    </div>

    <script>
        const chat = document.getElementById('chat');
        const prompt = document.getElementById('prompt');
        const send = document.getElementById('send');
        let messages = [];

        async function sendMessage() {
            const content = prompt.value.trim();
            if (!content) return;

            // Add user message
            messages.push({role: "user", content});
            addMsg("user", content);
            prompt.value = "";
            send.disabled = true;

            try {
                const resp = await fetch('/v1/chat/completions', {
                    method: 'POST',
                    headers: {'content-type': 'application/json'},
                    body: JSON.stringify({
                        messages,
                        max_tokens: 256,
                        helix: {
                            segments_budget: "budget:15%",
                            segments_salt: Math.random().toString(36).slice(2)
                        }
                    })
                });

                const data = await resp.json();
                const msg = data.choices[0].message.content;
                const helix = data.helix || {};
                const trace = data.trace || {};

                messages.push({role: "assistant", content: msg});
                addMsg("assistant", msg, {
                    model: trace.model_selected,
                    tag: trace.tag,
                    coverage: helix.selection?.decode_coverage_pct,
                    bitmask: helix.selection?.segments_mask_sha256?.slice(0,12)
                });

            } catch(e) {
                addMsg("assistant", "⚠️ Error: " + e.message, {error: true});
            }

            send.disabled = false;
            prompt.focus();
        }

        function addMsg(role, content, meta = {}) {
            const div = document.createElement('div');
            div.className = 'msg ' + role;
            div.textContent = content;

            if (Object.keys(meta).length > 0) {
                const metaDiv = document.createElement('div');
                metaDiv.className = 'meta';
                metaDiv.innerHTML = `
                    ${meta.tag ? `Tag: <b>${meta.tag}</b>` : ''}
                    ${meta.model ? `Model: <b>${meta.model}</b>` : ''}
                    ${meta.coverage ? `Coverage: ${meta.coverage}%` : ''}
                    ${meta.bitmask ? `Mask: ${meta.bitmask}...` : ''}
                `;
                div.appendChild(metaDiv);
            }

            chat.appendChild(div);
            chat.scrollTop = chat.scrollHeight;
        }

        prompt.addEventListener('keypress', e => {
            if (e.key === 'Enter' && !send.disabled) sendMessage();
        });
        send.addEventListener('click', sendMessage);
        prompt.focus();
    </script>
</body>
</html>
"""


@app.post("/v1/chat/completions")
async def chat_completions(req: ChatRequest):
    """OpenAI-compatible chat endpoint with model routing."""
    if not req.messages:
        raise HTTPException(400, "messages required")

    # Extract last user message for tagging
    last_msg = next((m for m in reversed(req.messages) if m.role == "user"), None)
    if not last_msg:
        raise HTTPException(400, "no user message found")

    # Tag and route
    tag = tag_prompt(last_msg.content)
    model = MODEL_ROUTES.get(tag, "flowtorch/braid")

    # Forward to FlowtTorch with routing
    async with httpx.AsyncClient() as client:
        try:
            ft_resp = await client.post(
                f"{FT_BASE}/v1/chat/completions",
                json={
                    "model": model,
                    "messages": [{"role": m.role, "content": m.content} for m in req.messages],
                    "max_tokens": req.max_tokens,
                    "helix": req.helix or {}
                },
                timeout=30.0
            )
            ft_resp.raise_for_status()
            data = ft_resp.json()

            # Add routing trace
            if "trace" not in data:
                data["trace"] = {}
            data["trace"]["tag"] = tag
            data["trace"]["model_selected"] = model
            data["trace"]["router"] = "mock_chatgpt"

            return data

        except httpx.HTTPError as e:
            raise HTTPException(500, f"FlowtTorch error: {e}")


@app.get("/health")
async def health():
    """Health check."""
    async with httpx.AsyncClient() as client:
        try:
            resp = await client.get(f"{FT_BASE}/health", timeout=5.0)
            ft_ok = resp.status_code == 200
        except:
            ft_ok = False

    return {
        "status": "ok" if ft_ok else "degraded",
        "flowtorch": "connected" if ft_ok else "unreachable",
        "router": "active"
    }


if __name__ == "__main__":
    import uvicorn
    port = int(os.environ.get("MOCK_PORT", "8090"))
    uvicorn.run(app, host="127.0.0.1", port=port, log_level="info")
