#!/usr/bin/env python3
"""
OpenAI-compatible facade in front of Ollama.

1. Injects a SearXNG-backed `web_search` tool into every chat/completions request
2. Executes web_search tool calls server-side (non-streaming requests) and feeds
   the results back to the model before answering the client
3. All other tool calls (the client's own tools) are returned to the client
   untouched — the client runs its own tool loop
4. Streaming requests are passed straight through to Ollama (no probe), so the
   client keeps native SSE and thinking output. In that case, if the model calls
   web_search the client receives that tool_call.

Ollama ignores the Authorization header (no API key), so clients may send any key.
"""
from flask import Flask, request, Response, stream_with_context
import requests
import json
import sys
import os

app = Flask(__name__)

sys.stdout.reconfigure(line_buffering=True)
sys.stderr.reconfigure(line_buffering=True)

BACKEND_URL = os.environ.get("BACKEND_URL", "http://vllm:8000")
BACKEND_API_KEY = os.environ.get("BACKEND_API_KEY", "")
SEARXNG_URL = os.environ.get("SEARXNG_URL", "http://host.docker.internal:8889")
SEARXNG_RESULTS = int(os.environ.get("SEARXNG_RESULTS", "5"))
MAX_SEARCH_ITER = int(os.environ.get("MAX_SEARCH_ITER", "10"))

SEARXNG_TOOL = {
    "type": "function",
    "function": {
        "name": "web_search",
        "description": (
            "Search the web for current information. Use this whenever you need "
            "up-to-date facts, recent events, or anything that may have changed "
            "since your training cutoff. Do not guess — call web_search instead."
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "The search query string"
                }
            },
            "required": ["query"]
        }
    }
}

SEARCH_SYSTEM_MSG = (
    "You have access to a web_search tool. "
    "Use it whenever the user asks about current events, real-time data, recent news, "
    "live prices, weather, sports scores, or anything else that may have changed since "
    "your training cutoff. Do not say you cannot access the internet — call web_search instead."
)


def searxng_search(query):
    """Execute a search against SearXNG and return a formatted result string."""
    try:
        print(f"[SEARXNG] Searching: {query}")
        resp = requests.get(
            f"{SEARXNG_URL}/search",
            params={"q": query, "format": "json", "categories": "general"},
            timeout=15
        )
        resp.raise_for_status()
        data = resp.json()
        results = data.get("results", [])[:SEARXNG_RESULTS]
        if not results:
            return f"No results found for: {query}"
        lines = [f"Search results for: {query}\n"]
        for i, r in enumerate(results, 1):
            title = r.get("title", "No title")
            url = r.get("url", "")
            content = r.get("content", "").strip()
            lines.append(f"{i}. {title}")
            if url:
                lines.append(f"   URL: {url}")
            if content:
                lines.append(f"   {content}")
            lines.append("")
        return "\n".join(lines).strip()
    except Exception as e:
        print(f"[ERROR] SearXNG search failed: {e}")
        return f"Search failed: {e}"


def inject_search(json_data):
    """Inject web_search tool and system prompt nudge if not already present."""
    tools = json_data.setdefault('tools', [])
    names = {t.get('function', {}).get('name') for t in tools if isinstance(t, dict)}
    if 'web_search' not in names:
        tools.append(SEARXNG_TOOL)
        print("[INJECTED] Added web_search tool")

    messages = json_data.setdefault('messages', [])
    already = any('web_search' in str(m.get('content') or '')
                  for m in messages if m.get('role') == 'system')
    if not already:
        if messages and messages[0].get('role') == 'system':
            # Append to the existing first system message so we stay at position 0
            messages[0]['content'] = str(messages[0]['content']) + '\n\n' + SEARCH_SYSTEM_MSG
            print("[INJECTED] Appended web_search to existing system prompt")
        else:
            # No system message yet — prepend one at position 0
            messages.insert(0, {"role": "system", "content": SEARCH_SYSTEM_MSG})
            print("[INJECTED] Added web_search system prompt at position 0")


def call_backend_nonstream(json_data, headers):
    """Call backend with streaming off, return parsed JSON or raise."""
    data = dict(json_data)
    data['stream'] = False
    data.pop('stream_options', None)
    body = json.dumps(data).encode('utf-8')
    hdrs = {k: v for k, v in headers.items() if k.lower() != 'host'}
    if BACKEND_API_KEY:
        hdrs['Authorization'] = f'Bearer {BACKEND_API_KEY}'
    hdrs['Content-Length'] = str(len(body))
    hdrs['Content-Type'] = 'application/json'
    resp = requests.post(f"{BACKEND_URL}/v1/chat/completions",
                         headers=hdrs, data=body, timeout=600)
    resp.raise_for_status()
    return resp.json()


def call_backend_passthrough(json_data, headers):
    """Forward request to backend as-is (preserving streaming), return raw Response."""
    body = json.dumps(json_data).encode('utf-8')
    hdrs = {k: v for k, v in headers.items() if k.lower() != 'host'}
    if BACKEND_API_KEY:
        hdrs['Authorization'] = f'Bearer {BACKEND_API_KEY}'
    hdrs['Content-Length'] = str(len(body))
    hdrs['Content-Type'] = 'application/json'
    sess = requests.Session()
    sess.headers.update(hdrs)
    resp = sess.post(f"{BACKEND_URL}/v1/chat/completions",
                     data=body, stream=True, timeout=(10, 600))

    def generate():
        count = 0
        try:
            for line in resp.iter_lines():
                if line:
                    count += 1
                    yield line + b'\n\n'
            print(f"[STREAM] Completed: {count} SSE lines forwarded")
        except Exception as e:
            print(f"[STREAM] ERROR after {count} lines: {e}")
            import traceback; traceback.print_exc()

    out_headers = {k: v for k, v in resp.headers.items()
                   if k.lower() in {'content-type', 'cache-control', 'x-accel-buffering'}}
    return Response(generate(), status=resp.status_code, headers=out_headers)


def run_loop(json_data, headers):
    """
    Non-streaming path. Calls the model, handles web_search calls transparently,
    returns to client when the model produces non-search output.

    If the model calls the client's own tools alongside or instead of
    web_search, those are returned to the client immediately.
    """
    is_streaming = json_data.get('stream', False)
    internal = dict(json_data)
    internal['messages'] = list(json_data.get('messages', []))

    searches = 0
    empty = 0
    for iteration in range(MAX_SEARCH_ITER):
        try:
            completion = call_backend_nonstream(internal, headers)
        except Exception as e:
            print(f"[ERROR] Backend call failed: {e}")
            import traceback; traceback.print_exc()
            return Response(
                json.dumps({'error': {'message': str(e), 'type': 'proxy_error'}}).encode(),
                status=502, headers={'Content-Type': 'application/json'}
            )

        choice = completion.get('choices', [{}])[0]
        message = choice.get('message', {})
        tool_calls = message.get('tool_calls') or []

        search_calls = [tc for tc in tool_calls if tc.get('function', {}).get('name') == 'web_search']
        client_calls = [tc for tc in tool_calls if tc.get('function', {}).get('name') != 'web_search']

        if not search_calls:
            print(f"[LOOP] Done after {iteration} search iteration(s)")
            if client_calls:
                print(f"[LOOP] Returning {len(client_calls)} client tool call(s) to client")
            if iteration == 0 and is_streaming:
                print("[LOOP] No search needed, passing through natively")
                return call_backend_passthrough(json_data, headers)
            if is_streaming:
                return _stream_completion(completion)
            return Response(json.dumps(completion).encode(),
                            status=200, headers={'Content-Type': 'application/json'})

        # Model wants to search — execute and loop
        print(f"[LOOP] Iteration {iteration}: {len(search_calls)} search call(s), "
              f"{len(client_calls)} client call(s)")

        asst = {"role": "assistant", "tool_calls": tool_calls}
        if message.get('content'):
            asst['content'] = message['content']
        internal['messages'].append(asst)

        for tc in search_calls:
            try:
                args = json.loads(tc.get('function', {}).get('arguments', '{}'))
            except Exception:
                args = {}
            result = searxng_search(args.get('query', ''))
            searches += 1
            if result.startswith(("No results found", "Search failed")):
                empty += 1
            print(f"[LOOP] Search returned {len(result)} chars")
            internal['messages'].append({
                "role": "tool",
                "tool_call_id": tc.get('id', 'search_result'),
                "content": result
            })

        # Client tool calls that arrive alongside a search cannot be executed
        # here — stub them so the model does not stall waiting for results
        for tc in client_calls:
            name = tc.get('function', {}).get('name', 'unknown')
            internal['messages'].append({
                "role": "tool",
                "tool_call_id": tc.get('id', 'stub'),
                "content": f"[{name} result will be provided by the user's environment]"
            })

    print(f"[ERROR] Exceeded max search iterations ({searches} searches, {empty} empty)")
    hint = " — SearXNG returned no usable results; check the SearXNG engines" if empty else ""
    return Response(
        json.dumps({'error': {
            'message': f'Exceeded max search iterations ({searches} searches, {empty} empty){hint}',
            'type': 'proxy_error'}}).encode(),
        status=500, headers={'Content-Type': 'application/json'}
    )


def _stream_completion(completion):
    """Re-stream a non-streaming completion as SSE."""
    choice = completion.get('choices', [{}])[0]
    message = choice.get('message', {})
    content = message.get('content') or ''
    reasoning = message.get('reasoning') or message.get('reasoning_content') or ''
    tool_calls = message.get('tool_calls') or []
    finish_reason = choice.get('finish_reason', 'stop')
    cid = completion.get('id', 'chatcmpl-search')
    model = completion.get('model', 'unknown')
    usage = completion.get('usage')

    def generate():
        yield _sse({'id': cid, 'object': 'chat.completion.chunk', 'model': model,
                    'choices': [{'index': 0, 'delta': {'role': 'assistant', 'content': ''},
                                 'finish_reason': None}]})

        if tool_calls:
            for i, tc in enumerate(tool_calls):
                yield _sse({'id': cid, 'object': 'chat.completion.chunk', 'model': model,
                            'choices': [{'index': 0, 'delta': {'tool_calls': [{
                                'index': i,
                                'id': tc.get('id', ''),
                                'type': 'function',
                                'function': {
                                    'name': tc.get('function', {}).get('name', ''),
                                    'arguments': tc.get('function', {}).get('arguments', '')
                                }
                            }]}, 'finish_reason': None}]})
        else:
            if reasoning:
                for i in range(0, len(reasoning), 64):
                    yield _sse({'id': cid, 'object': 'chat.completion.chunk', 'model': model,
                                'choices': [{'index': 0, 'delta': {'reasoning_content': reasoning[i:i + 64]},
                                             'finish_reason': None}]})
            for i in range(0, max(len(content), 1), 64):
                yield _sse({'id': cid, 'object': 'chat.completion.chunk', 'model': model,
                            'choices': [{'index': 0, 'delta': {'content': content[i:i + 64]},
                                         'finish_reason': None}]})

        final = {'id': cid, 'object': 'chat.completion.chunk', 'model': model,
                 'choices': [{'index': 0, 'delta': {}, 'finish_reason': finish_reason}]}
        if usage:
            final['usage'] = usage
        yield _sse(final)
        yield b"data: [DONE]\n\n"

    return Response(stream_with_context(generate()), status=200,
                    headers={'Content-Type': 'text/event-stream',
                             'Cache-Control': 'no-cache',
                             'X-Accel-Buffering': 'no'})


def _sse(obj):
    return f"data: {json.dumps(obj)}\n\n".encode()


def forward(method, url, headers, data):
    """Transparently forward a request and stream the response back."""
    hdrs = {k: v for k, v in headers.items() if k.lower() != 'host'}
    if BACKEND_API_KEY:
        hdrs['Authorization'] = f'Bearer {BACKEND_API_KEY}'
    sess = requests.Session()
    sess.headers.update(hdrs)
    resp = sess.request(method=method, url=url, data=data,
                        stream=True, timeout=(10, 600))

    def generate():
        for line in resp.iter_lines():
            if line:
                yield line + b'\n\n'

    out_headers = {k: v for k, v in resp.headers.items()
                   if k.lower() in {'content-type', 'cache-control', 'x-accel-buffering'}}
    return Response(generate(), status=resp.status_code, headers=out_headers)


@app.route('/<path:path>', methods=['GET', 'POST', 'PUT', 'DELETE', 'PATCH'])
def proxy(path):
    url = f"{BACKEND_URL}/{path}"
    headers = dict(request.headers)
    headers.pop('Host', None)
    data = request.get_data()

    if request.method == 'POST' and '/chat/completions' in path:
        try:
            json_data = json.loads(data)
            client_tools = json_data.get('tools', [])
            is_streaming = json_data.get('stream', False)
            print(f"[DEBUG] Request keys: {list(json_data.keys())}, "
                  f"tools={len(client_tools)}, stream={is_streaming}")

            inject_search(json_data)

            if is_streaming:
                print("[FASTPATH] Streaming request — skipping probe, passing through natively")
                return call_backend_passthrough(json_data, headers)

            return run_loop(json_data, headers)

        except Exception as e:
            print(f"[ERROR] {e}")
            import traceback; traceback.print_exc()
            return Response(
                json.dumps({'error': {'message': str(e), 'type': 'proxy_error'}}).encode(),
                status=500, headers={'Content-Type': 'application/json'}
            )

    # Non-chat requests — forward directly
    try:
        return forward(request.method, url, headers, data)
    except Exception as e:
        print(f"[ERROR] Request failed: {e}")
        return Response(
            json.dumps({'error': {'message': str(e)}}).encode(),
            status=500, headers={'Content-Type': 'application/json'}
        )


if __name__ == '__main__':
    print("Starting PASSTHROUGH+SEARCH proxy (Ollama backend)")
    print(f"Forwarding to: {BACKEND_URL}")
    print(f"SearXNG: {SEARXNG_URL} (max {SEARXNG_RESULTS} results per search)")
    app.run(host='0.0.0.0', port=11536, threaded=True)