"""Standalone SSE diagnostic using http.client and keep-alive.

Run: .venv/bin/python scripts/test_deepseek_proxy_stream.py
No OpenAI SDK path joining, retries, or response cache is used.
"""

from __future__ import annotations

import argparse
from http.client import HTTPConnection, HTTPSConnection, IncompleteRead
import json
import os
import re
import socket
import sys
import time
from pathlib import Path
from urllib.parse import parse_qsl, quote, quote_plus, urlsplit, urlunsplit

def load_env_values(path: Path) -> dict[str, str]:
    """Read just the three diagnostic settings, with no third-party dependency."""
    values = {}
    if not path.is_file():
        return values
    for line in path.read_text(encoding="utf-8-sig").splitlines():
        line = line.strip()
        if line.startswith("export "):
            line = line[7:]
        if not line or line.startswith("#") or "=" not in line:
            continue
        key, value = line.split("=", 1)
        key, value = key.strip(), value.strip()
        if key not in {"DEEPSEEK_BASE_URL", "DEEPSEEK_API_KEY", "DEEPSEEK_MODEL"}:
            continue
        if value.startswith(("'", '"')):
            quote_char = value[0]
            end = value.find(quote_char, 1)
            if end < 0 or (value[end + 1:].strip() and not value[end + 1:].strip().startswith("#")):
                raise ValueError("invalid_env_quoting")
            value = value[1:end]
        else:
            value = value.split(" #", 1)[0].strip()
        values[key] = value
    return values


def _redact(value: str, secrets: list[str]) -> str:
    for secret in sorted(set(secrets), key=len, reverse=True):
        if secret:
            for variant in (secret, quote(secret, safe=""), quote_plus(secret)):
                value = value.replace(variant, "[REDACTED]")
    value = re.sub(r"(?i)(proxyToken[\"'\s]*[:=][\"'\s]*)[^&\s\"'<>]+", r"\1[REDACTED]", value)
    return re.sub(r"(?i)(Bearer\s+)\S+", r"\1[REDACTED]", value)


def probe(*, timeout: float = 30, method: str = "POST") -> dict:
    """Return safe diagnostics; environment variables override the root .env."""
    if method != "POST":
        return {"ok": False, "code": "invalid_method"}
    env_file = Path(__file__).resolve().parents[1] / ".env"
    try:
        values = load_env_values(env_file)
    except (OSError, ValueError):
        return {"ok": False, "code": "invalid_env_file", "hint": ".env无法读取或引号格式无效"}
    values.update({key: os.environ[key] for key in (
        "DEEPSEEK_BASE_URL", "DEEPSEEK_API_KEY", "DEEPSEEK_MODEL"
    ) if key in os.environ})
    url = values.get("DEEPSEEK_BASE_URL", "").strip()
    api_key = values.get("DEEPSEEK_API_KEY", "").strip()
    model = values.get("DEEPSEEK_MODEL", "deepseek-flash").strip()
    result = {
        "ok": False,
        "env_file_exists": env_file.is_file(),
        "api_key_configured": bool(api_key),
        "environment_proxy_configured": any(os.environ.get(k) for k in (
            "HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY", "http_proxy", "https_proxy", "all_proxy"
        )),
        "use_environment_proxy": False,
        "transport": "http.client",
        "request_connection": "keep-alive",
        "stream": True,
    }
    secrets = [api_key]
    try:
        parts = urlsplit(url)
        params = parse_qsl(parts.query, keep_blank_values=True)
        secrets.extend(value for key, value in params if key.lower() == "proxytoken")
        if (parts.scheme not in ("http", "https") or not parts.hostname
                or parts.username or parts.password or parts.fragment
                or any(char.isspace() for char in url)):
            raise ValueError
        _ = parts.port
        if parts.path.rstrip("/").endswith("/chat/completions"):
            parts = parts._replace(path=parts.path.rstrip("/"))
            url = urlunsplit(parts)
    except ValueError:
        return {**result, "code": "invalid_base_url", "hint": "DEEPSEEK_BASE_URL须为单行完整HTTP(S)代理URL"}
    result.update(endpoint=f"{parts.scheme}://{parts.netloc}{parts.path}",
                  query_keys=sorted({key for key, _ in params}),
                  model=_redact(model, secrets), method=method,
                  chat_payload_sent=method == "POST")
    query = dict(params)
    if query.get("servertype") != "deepseek" or not query.get("proxyToken"):
        return {**result, "code": "missing_proxy_params", "hint": "URL须包含servertype=deepseek及非空proxyToken，注意参数大小写"}

    payload = {
        "model": model,
        "messages": [{"role": "user", "content": "1+1等于多少"}],
        "max_tokens": 128,
        "stream": True,
    }
    headers = {"Accept": "*/*", "Connection": "keep-alive"}
    if method == "POST":
        headers["Content-Type"] = "application/json"
    if api_key:
        headers["Authorization"] = f"Bearer {api_key}"
    # GET probes the documented gateway URL only. Do not guess how a custom
    # gateway encodes upstream chat requests in its query parameters.
    request_body = json.dumps(payload, ensure_ascii=False).encode("utf-8") if method == "POST" else None
    started = time.perf_counter()
    connection = None
    try:
        connection_type = HTTPSConnection if parts.scheme == "https" else HTTPConnection
        connection = connection_type(parts.hostname, parts.port, timeout=timeout)
        target = (parts.path or "/") + ("?" + parts.query if parts.query else "")
        # Direct connection: no environment proxy, redirects, or automatic retry.
        connection.request(method, target, body=request_body, headers=headers)
        response = connection.getresponse()
        with response:
            result["http_status"] = response.status
            result["http_ok"] = 200 <= response.status < 300
            result["content_type"] = _redact(response.headers.get("Content-Type", ""), secrets)
            result["transfer_encoding"] = _redact(response.headers.get("Transfer-Encoding", ""), secrets)[:120]
            length = response.headers.get("Content-Length", "")
            expected_length = int(length) if length.isascii() and length.isdecimal() else None
            result["content_length"] = expected_length
            if result['http_ok'] and result['content_type'].split(';')[0].strip().lower() == 'text/event-stream':
                read_sse(response, result, secrets, started, timeout)
                return result
            result["response_complete"] = True
            try:
                raw = response.read(1024 * 1024 + 1)
            except IncompleteRead as exc:
                # Malformed/unfinished chunked responses can contain the useful
                # gateway error text in partial; do not discard that evidence.
                raw = exc.partial
                result.update(response_complete=False, error_type="IncompleteRead")
            if (expected_length is not None and not result["transfer_encoding"]
                    and len(raw) < expected_length):
                result.update(response_complete=False, error_type="ContentLengthMismatch")
            result["received_bytes"] = len(raw)
        if len(raw) > 1024 * 1024:
            result["code"] = "response_too_large"
            return result
        body = raw.decode("utf-8", errors="replace")
        result["response_preview"] = _redact(body, secrets)[:1200]
        if "authentication fails" in body.lower() and "governor" in body.lower():
            result.update(code="proxy_authentication_failed", hint="代理已响应但鉴权失败，请核对参数大小写、token及代理权限")
        elif not 200 <= result["http_status"] < 300:
            result["code"] = "http_error"
        elif not result["response_complete"]:
            result["code"] = "response_incomplete"
        else:
            try:
                data = json.loads(body)
                content = data["choices"][0]["message"]["content"]
                if not isinstance(content, str) or not content.strip():
                    raise ValueError
                result.update(code="unexpected_response", hint="请求stream=true，但返回普通JSON；代理未返回SSE流")
            except (ValueError, KeyError, TypeError, IndexError):
                result.update(code="unexpected_response", hint="HTTP成功但不是含非空choices[0].message.content的聊天响应；需核对代理返回协议")
        if not result["response_complete"]:
            result["hint"] = (
                f"已收到HTTP {result['http_status']}响应头，但响应体传输不完整；"
                + ("未取得可解析的正文。" if not raw else "response_preview为已收到的脱敏片段。")
                + "HTTP 200不代表聊天成功，请代理同事检查重复Transfer-Encoding、分块编码及结束标记"
            )
    except (TimeoutError, socket.timeout):
        result["code"] = "timeout"
    except OSError as exc:
        result["code"] = "connection_failed"
        result["error_type"] = type(exc).__name__
    except Exception as exc:
        # Do not expose exceptions containing the complete URL or credentials.
        result.update(code="request_failed", error_type=type(exc).__name__)
    finally:
        if connection is not None:
            connection.close()
        result["elapsed_seconds"] = round(time.perf_counter() - started, 3)
    return result


def read_sse(response, result, secrets, started, timeout):
    """Parse bounded SSE events incrementally, then verify HTTP framing to EOF."""
    result.update(response_complete=False, done_received=False, finish_reason=None,
                  content_chunks=0, received_bytes=0, first_content_seconds=None,
                  last_content_seconds=None, content_arrival_seconds=[])
    pending = b''
    event = []
    content = []
    invalid = False
    provider_error = False

    def dispatch():
        nonlocal invalid, provider_error
        if not event:
            return
        value = b'\n'.join(event).decode('utf-8')
        event.clear()
        if result['done_received']:
            invalid = True
            return
        if value == '[DONE]':
            result['done_received'] = True
            return
        data = json.loads(value)
        if not isinstance(data, dict):
            invalid = True
            return
        if 'error' in data:
            provider_error = True
            return
        choices = data.get('choices')
        if choices == []:  # Optional usage frame.
            return
        if not isinstance(choices, list) or len(choices) != 1 or not isinstance(choices[0], dict):
            invalid = True
            return
        choice = choices[0]
        delta = choice.get('delta', {})
        if choice.get('index', 0) != 0 or not isinstance(delta, dict):
            invalid = True
            return
        text = delta.get('content')
        if text is not None and not isinstance(text, str):
            invalid = True
            return
        if text:
            if result['finish_reason'] is not None:
                invalid = True
            content.append(text)
            elapsed = round(time.perf_counter() - started, 3)
            result['content_chunks'] += 1
            if result['first_content_seconds'] is None:
                result['first_content_seconds'] = elapsed
            result['last_content_seconds'] = elapsed
            if len(result['content_arrival_seconds']) < 100:
                result['content_arrival_seconds'].append(elapsed)
        finish = choice.get('finish_reason')
        if finish is not None:
            if not isinstance(finish, str) or result['finish_reason'] is not None:
                invalid = True
            result['finish_reason'] = _redact(str(finish), secrets)[:80]

    try:
        while True:
            if time.perf_counter() - started > timeout:
                raise TimeoutError
            chunk = response.read1(4096)
            if not chunk:
                break
            result['received_bytes'] += len(chunk)
            if result['received_bytes'] > 1024 * 1024:
                result['code'] = 'response_too_large'
                return
            pending += chunk
            while b'\n' in pending:
                line, pending = pending.split(b'\n', 1)
                line = line.rstrip(b'\r')
                if not line:
                    dispatch()
                elif line.startswith(b'data:'):
                    value = line[5:]
                    event.append(value[1:] if value.startswith(b' ') else value)
        if result['content_length'] is not None and not result['transfer_encoding'] and result['received_bytes'] != result['content_length']:
            raise IncompleteRead(b'')
        result['response_complete'] = True
        if pending or event:
            invalid = True
        if provider_error:
            result['code'] = 'provider_error'
        elif invalid:
            result['code'] = 'invalid_sse'
        elif not result['done_received'] or result['finish_reason'] is None:
            result['code'] = 'stream_incomplete'
        elif result['finish_reason'] != 'stop':
            result['code'] = 'incomplete_completion'
        elif not ''.join(content).strip():
            result['code'] = 'empty_stream_content'
        else:
            result.update(ok=True, code='stream_success')
    except IncompleteRead:
        result.update(code='response_incomplete', error_type='IncompleteRead')
    except (ValueError, UnicodeError):
        result['code'] = 'invalid_sse'
    finally:
        result['assistant_text'] = _redact(''.join(content), secrets)[:1200]
        result['hint'] = '到达时间仅供检查缓冲；单次短回答不能证明代理始终实时透传。成功需完整HTTP响应、非空正文、finish_reason=stop及[DONE]'


def main() -> int:
    parser = argparse.ArgumentParser(description="读取.env，以keep-alive POST测试DeepSeek SSE流")
    parser.add_argument("--timeout", type=float, default=30, help="请求超时秒数，默认30")
    parser.add_argument("--method", choices=("POST",), default="POST")
    args = parser.parse_args()
    if not 0 < args.timeout <= 120:
        parser.error("--timeout须大于0且不超过120秒")
    result = probe(timeout=args.timeout, method=args.method)
    print(json.dumps(result, ensure_ascii=False, indent=2))
    return 0 if result["ok"] else 1


if __name__ == "__main__":
    sys.exit(main())