#!/usr/bin/env python3
"""Same usage as 影刀 RPA:

    import pymysql
    conn = pymysql.connect(
        host="127.0.0.1",
        port=13306,
        user="root",
        password="123456",
        db="your_database",
    )

Values come from config.yaml: listen + local_auth + remote.database.
mysql-tunnel.exe must already be running.

  python tests/http_tunnel.py
  python tests/http_tunnel.py --sql "SHOW TABLES"
"""

from __future__ import annotations

import argparse
import base64
import ssl
import struct
import sys
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any
from urllib.error import HTTPError, URLError
from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse
from urllib.request import Request, urlopen

MAGIC = 1111
MARK_LONG = 0xFE
MARK_NULL = 0xFF
SEP_MORE = 0x01
SEP_LAST = 0x00
MAX_BLOCK = 64 << 20

ROOT = Path(__file__).resolve().parents[1]


class TunnelError(Exception):
    pass


def parse_scalar(raw: str) -> Any:
    v = raw.strip()
    if v == "":
        return ""
    if (v.startswith('"') and v.endswith('"')) or (v.startswith("'") and v.endswith("'")):
        return v[1:-1]
    low = v.lower()
    if low in ("true", "yes"):
        return True
    if low in ("false", "no"):
        return False
    try:
        return int(v)
    except ValueError:
        return v


def load_project_yaml(path: Path) -> dict[str, Any]:
    root: dict[str, Any] = {}
    section: str | None = None
    text = path.read_text(encoding="utf-8")
    for line in text.splitlines():
        stripped = line.strip()
        if not stripped or stripped.startswith("#"):
            continue
        indent = len(line) - len(line.lstrip(" \t"))
        if indent == 0 and stripped.endswith(":") and ":" not in stripped[:-1]:
            section = stripped[:-1]
            root[section] = {}
            continue
        if section is None or ":" not in stripped:
            continue
        key, _, val = stripped.partition(":")
        root[section][key.strip()] = parse_scalar(val)
    return root


def redact_url(raw: str) -> str:
    u = urlparse(raw)
    q = []
    for k, v in parse_qsl(u.query, keep_blank_values=True):
        if k.lower() in ("token", "password", "pass", "access_token", "accesstoken"):
            q.append((k, "***"))
        else:
            q.append((k, v))
    return urlunparse(u._replace(query=urlencode(q), fragment=""))


def looks_html(body: bytes) -> bool:
    s = body.lstrip()[:64].decode("latin-1", errors="replace").lower()
    return s.startswith("<!doctype") or s.startswith("<html")


@dataclass
class Header:
    magic: int
    version: int
    errno: int
    message: str = ""

    @property
    def ok(self) -> bool:
        return self.errno == 0


@dataclass
class Field:
    name: str
    table: str
    type: int = 0
    flags: int = 0
    length: int = 0


@dataclass
class Result:
    errno: int = 0
    message: str = ""
    affected_rows: int = 0
    insert_id: int = 0
    num_fields: int = 0
    num_rows: int = 0
    info: str = ""
    fields: list[Field] = field(default_factory=list)
    rows: list[list[str | None]] = field(default_factory=list)

    @property
    def is_error(self) -> bool:
        return self.errno > 0

    @property
    def is_result_set(self) -> bool:
        return self.errno == 0 and self.num_fields > 0


class Reader:
    def __init__(self, buf: bytes) -> None:
        self.buf = buf
        self.off = 0

    def remaining(self) -> int:
        return len(self.buf) - self.off

    def read(self, n: int) -> bytes:
        if n < 0 or self.off + n > len(self.buf):
            raise TunnelError("truncated Navicat binary")
        s = self.buf[self.off : self.off + n]
        self.off += n
        return s

    def u8(self) -> int:
        return self.read(1)[0]

    def u16(self) -> int:
        return struct.unpack(">H", self.read(2))[0]

    def u32(self) -> int:
        return struct.unpack(">I", self.read(4))[0]

    def block(self, nullable: bool = False) -> tuple[bytes | None, bool]:
        if self.remaining() == 0:
            raise TunnelError("truncated GetBlock")
        first = self.buf[self.off]
        if nullable and first == MARK_NULL:
            self.off += 1
            return None, True
        self.off += 1
        if first < 254:
            n = first
        elif first == MARK_LONG:
            n = self.u32()
            if n > MAX_BLOCK:
                raise TunnelError(f"block length {n} exceeds 64MiB")
        else:
            raise TunnelError(f"invalid block marker 0x{first:02x}")
        return self.read(n), False

    def block_str(self) -> str:
        val, is_null = self.block(False)
        if is_null or val is None:
            raise TunnelError("unexpected NULL block")
        return val.decode("utf-8", errors="replace")


def parse_header(body: bytes) -> tuple[Header, Reader]:
    if looks_html(body):
        raise TunnelError("tunnel returned HTML, not Navicat binary")
    if len(body) < 16:
        raise TunnelError(f"response too short ({len(body)} bytes)")
    r = Reader(body)
    magic = r.u32()
    if magic != MAGIC:
        raise TunnelError(f"invalid Navicat magic {magic} (want {MAGIC})")
    version = r.u16()
    errno = r.u32()
    r.read(6)
    msg = ""
    if errno != 0:
        msg = r.block_str()
    return Header(magic, version, errno, msg), r


def parse_connect(body: bytes) -> tuple[Header, dict[str, str]]:
    h, r = parse_header(body)
    if not h.ok:
        return h, {}
    return h, {
        "host_info": r.block_str(),
        "proto_info": r.block_str(),
        "server_info": r.block_str(),
    }


def parse_one_result(r: Reader) -> Result:
    res = Result(
        errno=r.u32(),
        affected_rows=r.u32(),
        insert_id=r.u32(),
        num_fields=r.u32(),
        num_rows=r.u32(),
    )
    r.read(12)
    if res.errno > 0:
        res.message = r.block_str()
        return res
    if res.num_fields == 0:
        res.info = r.block_str()
        return res
    for _ in range(res.num_fields):
        res.fields.append(Field(r.block_str(), r.block_str(), r.u32(), r.u32(), r.u32()))
    for _ in range(res.num_rows):
        row: list[str | None] = []
        for _c in range(res.num_fields):
            val, is_null = r.block(True)
            row.append(None if is_null else (val or b"").decode("utf-8", errors="replace"))
        res.rows.append(row)
    return res


def parse_query(body: bytes) -> tuple[Header, list[Result]]:
    h, r = parse_header(body)
    if not h.ok:
        return h, []
    results: list[Result] = []
    while r.remaining() > 0:
        results.append(parse_one_result(r))
        if r.remaining() == 0:
            break
        sep = r.u8()
        if sep == SEP_LAST:
            break
        if sep != SEP_MORE:
            raise TunnelError(f"invalid query separator 0x{sep:02x}")
    return h, results


@dataclass
class Config:
    path: Path
    listen_host: str
    listen_port: int
    local_user: str
    local_password: str
    url: str
    encode_base64: bool
    tls_insecure: bool
    remote_host: str
    remote_port: int
    remote_user: str
    remote_password: str
    database: str
    charset: str
    timeout: int

    @classmethod
    def load(cls, path: Path) -> Config:
        data = load_project_yaml(path)
        listen = data.get("listen") or {}
        local = data.get("local_auth") or {}
        tunnel = data.get("tunnel") or {}
        remote = data.get("remote") or {}
        options = data.get("options") or {}
        url = str(tunnel.get("url") or "").strip()
        if not url:
            raise TunnelError(f"{path}: tunnel.url is required")
        local_user = str(local.get("username") or "").strip()
        if not local_user:
            raise TunnelError(f"{path}: local_auth.username is required")
        timeout = int(options.get("request_timeout") or 30)
        if timeout <= 0:
            timeout = 30
        return cls(
            path=path,
            listen_host=str(listen.get("host") or "127.0.0.1"),
            listen_port=int(listen.get("port") or 13306),
            local_user=local_user,
            local_password=str(local.get("password") or ""),
            url=url,
            encode_base64=bool(tunnel.get("encode_base64")),
            tls_insecure=bool(tunnel.get("tls_insecure_skip_verify")),
            remote_host=str(remote.get("host") or ""),
            remote_port=int(remote.get("port") or 3306),
            remote_user=str(remote.get("username") or ""),
            remote_password=str(remote.get("password") or ""),
            database=str(remote.get("database") or ""),
            charset=str(remote.get("charset") or "utf8mb4"),
            timeout=timeout,
        )


class HTTPClient:
    """POST actn=C/Q to ntunnel_mysql.php using remote.* (not local_auth)."""

    def __init__(self, cfg: Config) -> None:
        self.cfg = cfg
        ctx = ssl.create_default_context()
        if cfg.tls_insecure:
            ctx.check_hostname = False
            ctx.verify_mode = ssl.CERT_NONE
        self._ssl = ctx

    def post(self, actn: str, queries: list[str] | None = None) -> bytes:
        fields: list[tuple[str, str]] = [
            ("actn", actn),
            ("host", self.cfg.remote_host),
            ("port", str(self.cfg.remote_port) if self.cfg.remote_port > 0 else ""),
            ("login", self.cfg.remote_user),
            ("password", self.cfg.remote_password),
            ("db", self.cfg.database),
        ]
        qs = [q for q in (queries or []) if q.strip()]
        if self.cfg.encode_base64 and qs:
            fields.append(("encodeBase64", "1"))
            for q in qs:
                fields.append(("q[]", base64.b64encode(q.encode("utf-8")).decode("ascii")))
        else:
            for q in qs:
                fields.append(("q[]", q))
        body = urlencode(fields).encode("utf-8")
        req = Request(
            self.cfg.url,
            data=body,
            method="POST",
            headers={
                "Content-Type": "application/x-www-form-urlencoded",
                "User-Agent": "mysql-tunnel-python/0.5",
            },
        )
        kw: dict[str, Any] = {"timeout": self.cfg.timeout}
        if urlparse(self.cfg.url).scheme == "https":
            kw["context"] = self._ssl
        try:
            with urlopen(req, **kw) as resp:
                status = resp.status
                data = resp.read()
        except HTTPError as e:
            if e.code == 403:
                raise TunnelError("HTTP 403: invalid or missing token") from None
            raise TunnelError(f"HTTP {e.code}") from None
        except URLError as e:
            raise TunnelError(f"HTTP transport: {e.reason}") from None
        if status < 200 or status >= 300:
            raise TunnelError(f"HTTP {status}")
        if looks_html(data):
            raise TunnelError("tunnel returned HTML, not Navicat binary")
        return data

    def connect(self) -> tuple[Header, dict[str, str]]:
        return parse_connect(self.post("C"))

    def query(self, *sqls: str) -> tuple[Header, list[Result]]:
        qs: list[str] = []
        cs = (self.cfg.charset or "").strip()
        if cs and cs.lower() not in ("utf8", "utf8mb3") and cs.replace("_", "").isalnum():
            qs.append("SET NAMES " + cs)
        qs.extend(sqls)
        return parse_query(self.post("Q", qs))


def mysql_connect(cfg: Config):
    try:
        import pymysql
    except ImportError as e:
        raise TunnelError("install PyMySQL:  pip install pymysql") from e
    import warnings

    # Same signature as 影刀: host/port/user/password/db → listen + local_auth + remote.database
    try:
        with warnings.catch_warnings():
            warnings.simplefilter("ignore", DeprecationWarning)
            return pymysql.connect(
                host=cfg.listen_host,
                port=cfg.listen_port,
                user=cfg.local_user,
                password=cfg.local_password,
                db=cfg.database,
            )
    except pymysql.err.OperationalError as e:
        code = e.args[0] if e.args else 0
        if code in (2003, 2002):
            raise TunnelError(
                f"cannot connect {cfg.listen_host}:{cfg.listen_port} — "
                "start mysql-tunnel.exe first (GUI 「启动」 or --console)"
            ) from None
        if code == 1045:
            raise TunnelError(
                f"local_auth rejected ({cfg.local_user}) — "
                "影刀/本脚本的 user/password 必须等于 config.yaml 的 local_auth"
            ) from None
        raise TunnelError(f"mysql connect: {e}") from None


def query_mysql(conn, sql: str) -> Result:
    import pymysql

    cur = conn.cursor()
    try:
        cur.execute(sql)
        rows = cur.fetchall()
        desc = cur.description or ()
        res = Result(
            affected_rows=cur.rowcount if not desc else 0,
            insert_id=int(getattr(cur, "lastrowid", 0) or 0),
            num_fields=len(desc),
            num_rows=len(rows) if desc else 0,
        )
        if desc:
            res.fields = [Field(str(col[0]), "") for col in desc]
            for row in rows:
                out: list[str | None] = []
                for v in row:
                    if v is None:
                        out.append(None)
                    elif isinstance(v, bytes):
                        out.append(v.decode("utf-8", errors="replace"))
                    else:
                        out.append(str(v))
                res.rows.append(out)
        return res
    except pymysql.MySQLError as e:
        errno = int(e.args[0]) if e.args else 0
        msg = str(e.args[1]) if len(e.args) > 1 else str(e)
        return Result(errno=errno, message=msg)
    finally:
        cur.close()


def find_config(explicit: str | None) -> Path:
    cands: list[Path] = []
    if explicit:
        cands.append(Path(explicit))
    cands.extend(
        [
            ROOT / "config.yaml",
            Path.cwd() / "config.yaml",
            ROOT / "dist" / "config.yaml",
        ]
    )
    for p in cands:
        if p.is_file():
            return p.resolve()
    raise TunnelError("config.yaml not found; pass --config")


def last_result_set(results: list[Result]) -> Result | None:
    for res in reversed(results):
        if res.is_result_set:
            return res
    return results[-1] if results else None


def print_header(title: str) -> None:
    print()
    print("== " + title)


def print_result(res: Result, max_rows: int = 40) -> None:
    if res.is_error:
        print(f"  mysql errno {res.errno}: {res.message}")
        return
    if not res.is_result_set:
        extra = f"  info={res.info!r}" if res.info else ""
        print(f"  OK  affected={res.affected_rows}  insert_id={res.insert_id}{extra}")
        return
    names = [f.name or "?" for f in res.fields]
    print(f"  {res.num_rows} row(s), columns: {', '.join(names)}")
    widths = [len(n) for n in names]
    shown = res.rows[:max_rows]
    for row in shown:
        for i, cell in enumerate(row):
            text = "NULL" if cell is None else cell.replace("\n", "\\n")
            widths[i] = max(widths[i], min(len(text), 48))

    def fmt(val: str, w: int) -> str:
        if len(val) > w:
            val = val[: w - 1] + "…"
        return val.ljust(w)

    print("  " + " | ".join(fmt(n, widths[i]) for i, n in enumerate(names)))
    print("  " + "-+-".join("-" * w for w in widths))
    for row in shown:
        cells = []
        for i, cell in enumerate(row):
            text = "NULL" if cell is None else cell.replace("\n", "\\n")
            cells.append(fmt(text, widths[i]))
        print("  " + " | ".join(cells))
    if res.num_rows > max_rows:
        print(f"  … {res.num_rows - max_rows} more row(s)")


def default_sql() -> list[str]:
    return [
        "SELECT 1 AS ok",
        "SELECT VERSION() AS version, DATABASE() AS db, USER() AS user, "
        "@@character_set_connection AS charset",
        "SHOW TABLES",
    ]


def run_local(cfg: Config, extra_sql: list[str]) -> int:
    print("config:", cfg.path)
    print("pymysql.connect(")
    print(f'    host="{cfg.listen_host}",')
    print(f"    port={cfg.listen_port},")
    print(f'    user="{cfg.local_user}",')
    print('    password="***",')
    print(f'    db="{cfg.database}",')
    print(")")
    print("via:    mysql-tunnel.exe →", redact_url(cfg.url))

    try:
        conn = mysql_connect(cfg)
    except TunnelError as e:
        print()
        print("  Connect: FAIL")
        print(" ", e)
        return 1

    print()
    print("MySQL数据库连接成功")
    print("  Connect: OK  (local_auth)")
    code = 0
    try:
        for sql in extra_sql or default_sql():
            print_header(sql.replace("\n", " "))
            res = query_mysql(conn, sql)
            print_result(res)
            if res.is_error:
                code = 1
    finally:
        conn.close()
    return code


def run_direct(cfg: Config, extra_sql: list[str]) -> int:
    print("config:", cfg.path)
    print("mode:  --direct  (POST ntunnel, remote.* not local_auth)")
    print("tunnel:", redact_url(cfg.url))
    print("remote:", f"{cfg.remote_user}@{cfg.remote_host}:{cfg.remote_port}/{cfg.database}")

    client = HTTPClient(cfg)
    code = 0
    print_header("actn=C  connect")
    try:
        h, conninfo = client.connect()
    except TunnelError as e:
        print("  Tunnel: FAIL")
        print(" ", e)
        return 1
    if not h.ok:
        print(f"  errno {h.errno}: {h.message}")
        print("  Remote MySQL: FAIL")
        return 1
    print(f"  Tunnel: OK  (script version {h.version})")
    print("  Remote MySQL: OK")
    print(f"  server_info: {conninfo.get('server_info', '')}")

    for sql in extra_sql or default_sql():
        print_header("actn=Q  " + sql.replace("\n", " "))
        try:
            h, results = client.query(sql)
        except TunnelError as e:
            print("  FAIL", e)
            code = 1
            continue
        if not h.ok:
            print(f"  header errno {h.errno}: {h.message}")
            code = 1
            continue
        res = last_result_set(results)
        if res is None:
            print("  no result")
            code = 1
            continue
        print_result(res)
        if res.is_error:
            code = 1
    return code


def main(argv: list[str]) -> int:
    try:
        sys.stdout.reconfigure(encoding="utf-8")
    except Exception:
        pass
    ap = argparse.ArgumentParser(
        description="Connect through mysql-tunnel using local_auth (same as Yingdao)"
    )
    ap.add_argument("--config", default="", help="path to config.yaml")
    ap.add_argument("--sql", action="append", default=[], help="SQL to run (repeatable)")
    ap.add_argument(
        "--direct",
        action="store_true",
        help="POST ntunnel_mysql.php with remote.* instead of local_auth",
    )
    args = ap.parse_args(argv)
    try:
        cfg = Config.load(find_config(args.config or None))
        if args.direct:
            return run_direct(cfg, args.sql)
        return run_local(cfg, args.sql)
    except TunnelError as e:
        print("error:", e, file=sys.stderr)
        return 1


if __name__ == "__main__":
    raise SystemExit(main(sys.argv[1:]))