from flask import Flask, request, Response, stream_with_context, send_from_directory
import requests
import json
import os
import traceback
import sys

app = Flask(__name__)
STATIC_DIR = os.path.dirname(os.path.abspath(__file__))

# Import huggingface_hub if available (pre-installed on HF Spaces)
try:
    from huggingface_hub import InferenceClient
    HAS_HF_HUB = True
except ImportError:
    HAS_HF_HUB = False

@app.route('/')
def index():
    return send_from_directory(STATIC_DIR, 'index.html')

@app.route('/api/proxy', methods=['POST', 'OPTIONS'])
def proxy():
    if request.method == 'OPTIONS':
        return Response(headers={'Access-Control-Allow-Origin': '*'})

    data = request.get_json(force=True)
    url = data.get('url')
    if not url:
        return {'error': 'Missing url'}, 400

    method = data.get('method', 'POST').upper()
    headers = data.get('headers', {})
    filtered_headers = {}
    skip = {'origin', 'referer', 'host', 'cookie', 'set-cookie', 'cf-connecting-ip', 'cf-ray', 'cf-worker', 'x-forwarded-for', 'x-forwarded-proto', 'x-real-ip'}
    for k, v in headers.items():
        if k.lower() not in skip:
            filtered_headers[k] = v

    body = data.get('body')

    req_kwargs = {
        'method': method,
        'url': url,
        'headers': filtered_headers,
        'stream': True,
        'timeout': 180,
    }

    if body is not None:
        if isinstance(body, dict):
            req_kwargs['json'] = body
        elif isinstance(body, str):
            req_kwargs['data'] = body
            content_type_key = next((k for k in filtered_headers if k.lower() == 'content-type'), None)
            if not content_type_key:
                req_kwargs['headers']['Content-Type'] = 'application/json'
        else:
            req_kwargs['data'] = str(body)

    try:
        resp = requests.request(**req_kwargs)
    except requests.exceptions.Timeout:
        return {'error': 'Backend request timed out'}, 504
    except requests.exceptions.ConnectionError as e:
        return {'error': f'Connection error: {e}'}, 502
    except Exception as e:
        tb = traceback.format_exc()
        print(f'Proxy error for {url}: {e}\n{tb}', file=sys.stderr, flush=True)
        return {'error': f'{type(e).__name__}: {e}'}, 500

    excluded = {'transfer-encoding', 'connection', 'keep-alive', 'content-encoding', 'content-length'}
    proxy_headers = {k: v for k, v in resp.headers.items() if k.lower() not in excluded}

    is_sse = 'text/event-stream' in resp.headers.get('Content-Type', '').lower()

    if is_sse:
        def generate():
            try:
                for chunk in resp.iter_content(chunk_size=None, decode_unicode=False):
                    if chunk:
                        yield chunk
            except GeneratorExit:
                resp.close()
        return Response(
            stream_with_context(generate()),
            status=resp.status_code,
            headers={**proxy_headers, 'Cache-Control': 'no-cache', 'X-Accel-Buffering': 'no'}
        )
    else:
        return Response(
            resp.content,
            status=resp.status_code,
            headers=proxy_headers,
        )


@app.route('/api/proxy/hf', methods=['POST', 'OPTIONS'])
def proxy_hf():
    """Proxy HuggingFace Inference API via huggingface_hub internal routing."""
    if request.method == 'OPTIONS':
        return Response(headers={'Access-Control-Allow-Origin': '*'})

    data = request.get_json(force=True)
    model = data.get('model', '')
    body = data.get('body', {})
    token = data.get('token', '')
    is_chat = data.get('is_chat', False)
    stream = data.get('stream', False)

    if not model:
        return {'error': 'Missing model'}, 400

    if not HAS_HF_HUB:
        return {'error': 'huggingface_hub not available on server'}, 500

    try:
        client = InferenceClient(token=token)

        if is_chat and not stream:
            messages = body.get('messages', [])
            temperature = body.get('temperature', 0.7)
            max_tokens = body.get('max_tokens', 2048)

            result = client.chat.completions.create(
                model=model,
                messages=messages,
                temperature=temperature,
                max_tokens=max_tokens,
                stream=False,
            )
            return Response(
                json.dumps(result.to_dict()),
                status=200,
                headers={'Content-Type': 'application/json'}
            )

        elif stream:
            messages = body.get('messages', [])
            temperature = body.get('temperature', 0.7)
            max_tokens = body.get('max_tokens', 2048)

            result = client.chat.completions.create(
                model=model,
                messages=messages,
                temperature=temperature,
                max_tokens=max_tokens,
                stream=True,
            )

            def generate():
                try:
                    for chunk in result:
                        yield f'data: {json.dumps(chunk.to_dict())}\n\n'
                    yield 'data: [DONE]\n\n'
                except Exception as e:
                    print(f'HF stream error: {e}', file=sys.stderr, flush=True)
            return Response(
                stream_with_context(generate()),
                status=200,
                headers={'Content-Type': 'text/event-stream', 'Cache-Control': 'no-cache', 'X-Accel-Buffering': 'no'}
            )

        else:
            # Text generation via _post to bypass model metadata lookup
            resp = client._post(
                path=f"/models/{model}",
                json=body,
                stream=False,
            )
            return Response(
                resp.content,
                status=resp.status_code,
                headers={k: v for k, v in resp.headers.items() if k.lower() not in {'content-encoding', 'transfer-encoding', 'content-length'}},
                content_type=resp.headers.get('Content-Type', 'application/json'),
            )

    except ImportError:
        return {'error': 'huggingface_hub library not found'}, 500
    except Exception as e:
        tb = traceback.format_exc()
        print(f'HF proxy error for {model}: {e}\n{tb}', file=sys.stderr, flush=True)
        error_msg = f'{type(e).__name__}: {e}'
        return {'error': error_msg}, 500


if __name__ == '__main__':
    port = int(os.environ.get('PORT', 7860))
    app.run(host='0.0.0.0', port=port, debug=False)