# -*- coding: utf-8 -*-
"""
CloudDrive2 gRPC Proxy - Runs gRPC calls in an isolated process.

Kodi uses Python sub-interpreters that prevent grpcio's Cython extension
from being loaded more than once per process. This proxy runs as a
standalone Python process, imports grpc once, and exposes gRPC calls
via a simple HTTP+JSON API on localhost:19799.

Start:  python3 cd2_grpc_proxy.py
Stop:   Send POST /shutdown or kill the process.
"""

import sys
import os
import json
import threading
import queue
import uuid
import traceback
from http.server import HTTPServer, BaseHTTPRequestHandler

# Ensure vendor path
_vendor_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "vendors")
if os.path.isdir(_vendor_dir) and _vendor_dir not in sys.path:
    sys.path.insert(0, _vendor_dir)

import grpc
from google.protobuf.empty_pb2 import Empty
from google.protobuf.json_format import MessageToDict
import clouddrive_pb2 as pb
import clouddrive_pb2_grpc as pb_grpc

HOST = "127.0.0.1"
PORT = 19799
CD2_HOST = "localhost"
CD2_PORT = 19798

_token = None
_channel = None
_stub = None

# Active QR login sessions: session_id -> {"queue": Queue, "thread": Thread}
_qr_sessions = {}


def _get_token_path():
    """Get the token file path, matching cd2_utils convention."""
    # On CoreELEC/Kodi, the addon profile is at:
    #   /storage/.kodi/userdata/addon_data/plugin.video.clouddrive2/
    # Outside Kodi, use ~/.config/clouddrive2-plugin/
    kodi_path = os.path.expanduser(
        "~/.kodi/userdata/addon_data/plugin.video.clouddrive2/token.json")
    if os.path.exists(os.path.dirname(kodi_path)):
        return kodi_path
    alt_path = os.path.expanduser("~/.config/clouddrive2-plugin/token.json")
    return alt_path


def _load_token():
    """Load saved JWT token."""
    try:
        token_path = _get_token_path()
        if os.path.exists(token_path):
            with open(token_path, "r") as f:
                data = json.load(f)
                return data.get("token")
    except Exception:
        pass
    return None


def _save_token(token):
    """Save JWT token."""
    try:
        token_path = _get_token_path()
        token_dir = os.path.dirname(token_path)
        os.makedirs(token_dir, exist_ok=True)
        with open(token_path, "w") as f:
            json.dump({"token": token, "saved_at": 0}, f)
    except Exception:
        pass


def _ensure_connected():
    """Ensure gRPC channel and stub are initialized."""
    global _channel, _stub, _token
    if _channel is None:
        _channel = grpc.insecure_channel(f"{CD2_HOST}:{CD2_PORT}")
        _stub = pb_grpc.CloudDriveFileSrvStub(_channel)
        _token = _load_token()


def _auth_metadata():
    if _token:
        return [("authorization", f"Bearer {_token}")]
    return None


def _proto_to_dict(msg):
    """Convert a protobuf message to a JSON-serializable dict."""
    if msg is None:
        return None
    return MessageToDict(msg, preserving_proto_field_name=True,
                         always_print_fields_with_no_presence=True,
                         use_integers_for_enums=True)


def _handle_call(rpc_func, request, timeout=10):
    """Call an authenticated RPC and return a dict response."""
    _ensure_connected()
    metadata = _auth_metadata()
    try:
        resp = rpc_func(request, timeout=timeout, metadata=metadata)
        return {"success": True, "data": _proto_to_dict(resp)}
    except Exception as e:
        return {"success": False, "error": str(e)}


def _handle_stream(rpc_func, request, timeout=30):
    """Call a server-streaming RPC, collect all responses."""
    _ensure_connected()
    metadata = _auth_metadata()
    try:
        stream = rpc_func(request, timeout=timeout, metadata=metadata)
        results = []
        for msg in stream:
            results.append(_proto_to_dict(msg))
        return {"success": True, "data": results}
    except Exception as e:
        return {"success": False, "error": str(e)}


# ---- Request Handlers ----

def _sys_info(body):
    _ensure_connected()
    try:
        resp = _stub.GetSystemInfo(Empty(), timeout=5)
        return {"success": True, "data": _proto_to_dict(resp)}
    except Exception as e:
        return {"success": False, "error": str(e)}


def _register(body):
    _ensure_connected()
    try:
        req = pb.UserRegisterRequest()
        req.userName = body["userName"]
        req.password = body["password"]
        resp = _stub.Register(req, timeout=10)
        return {"success": True, "data": _proto_to_dict(resp)}
    except Exception as e:
        return {"success": False, "error": str(e)}


def _login(body):
    _ensure_connected()
    try:
        req = pb.UserLoginRequest()
        req.userName = body["userName"]
        req.password = body["password"]
        resp = _stub.Login(req, timeout=10)
        return {"success": True, "data": _proto_to_dict(resp)}
    except Exception as e:
        return {"success": False, "error": str(e)}


def _get_token(body):
    global _token
    _ensure_connected()
    try:
        req = pb.GetTokenRequest()
        req.userName = body["userName"]
        req.password = body["password"]
        resp = _stub.GetToken(req, timeout=10)
        if resp.success:
            _token = resp.token
            _save_token(resp.token)
        return {"success": True, "data": _proto_to_dict(resp)}
    except Exception as e:
        return {"success": False, "error": str(e)}


def _logout(body):
    """Logout from CloudFS server.
    Set logoutFromCloudFS=True to fully exit (close files, stop sync).
    Default True when body doesn't specify."""
    global _token
    _ensure_connected()
    metadata = _auth_metadata()
    try:
        req = pb.UserLogoutRequest()
        req.logoutFromCloudFS = body.get("logoutFromCloudFS", True)
        resp = _stub.Logout(req, timeout=10, metadata=metadata)
        _token = None
        return {"success": True, "data": _proto_to_dict(resp)}
    except Exception as e:
        return {"success": False, "error": str(e)}


def _runtime_info(body):
    return _handle_call(_stub.GetRuntimeInfo, Empty())


def _running_info(body):
    return _handle_call(_stub.GetRunningInfo, Empty())


def _system_settings(body):
    return _handle_call(_stub.GetSystemSettings, Empty())


def _set_system_settings(body):
    _ensure_connected()
    metadata = _auth_metadata()
    try:
        # body is a dict with systemsettings fields
        data = body.get("settings", body)
        req = pb.SystemSettings()
        for k, v in data.items():
            if hasattr(req, k):
                setattr(req, k, v)
        resp = _stub.SetSystemSettings(req, timeout=10, metadata=metadata)
        return {"success": True, "data": _proto_to_dict(resp)}
    except Exception as e:
        return {"success": False, "error": str(e)}


def _get_all_cloud_apis(body):
    return _handle_call(_stub.GetAllCloudApis, Empty())


def _get_cloud_memberships(body):
    return _handle_call(_stub.GetCloudMemberships, Empty())


def _can_add_more(body):
    return _handle_call(_stub.CanAddMoreCloudApis, Empty())


def _get_sub_files(body):
    path = body.get("path", "/")
    force_refresh = body.get("forceRefresh", False)
    _ensure_connected()
    metadata = _auth_metadata()
    try:
        req = pb.ListSubFileRequest()
        req.path = path
        req.forceRefresh = force_refresh
        stream = _stub.GetSubFiles(req, timeout=30, metadata=metadata)
        files = []
        for resp in stream:
            if resp.subFiles:
                for f in resp.subFiles:
                    files.append(_proto_to_dict(f))
        return {"success": True, "data": files}
    except Exception as e:
        return {"success": False, "error": str(e)}


def _get_download_url(body):
    path = body.get("path", "")
    get_direct = body.get("getDirectUrl", True)
    _ensure_connected()
    metadata = _auth_metadata()
    try:
        req = pb.GetDownloadUrlPathRequest()
        req.path = path
        req.get_direct_url = get_direct
        resp = _stub.GetDownloadUrlPath(req, timeout=10, metadata=metadata)
        return {"success": True, "data": _proto_to_dict(resp)}
    except Exception as e:
        return {"success": False, "error": str(e)}


def _get_mount_points(body):
    return _handle_call(_stub.GetMountPoints, Empty())


def _add_mount_point(body):
    _ensure_connected()
    metadata = _auth_metadata()
    try:
        req = pb.MountOption()
        req.mountPoint = body.get("mountPoint", body.get("cloudName", ""))
        req.sourceDir = body.get("sourceDir", "")
        req.localMount = body.get("localMount", True)
        req.autoMount = body.get("autoMount", True)
        req.readOnly = body.get("readOnly", False)
        req.uid = body.get("uid", 0)
        req.gid = body.get("gid", 0)
        req.permissions = body.get("permissions", "0755")
        req.name = body.get("name", body.get("cloudName", ""))
        resp = _stub.AddMountPoint(req, timeout=10, metadata=metadata)
        return {"success": True, "data": _proto_to_dict(resp)}
    except Exception as e:
        return {"success": False, "error": str(e)}


def _remove_mount_point(body):
    req = pb.MountPointRequest()
    req.MountPoint = body.get("cloudName", "")
    return _handle_call(_stub.RemoveMountPoint, req)


def _mount(body):
    req = pb.MountPointRequest()
    req.MountPoint = body.get("cloudName", "")
    return _handle_call(_stub.Mount, req)


def _unmount(body):
    req = pb.MountPointRequest()
    req.MountPoint = body.get("cloudName", "")
    return _handle_call(_stub.Unmount, req)


def _remove_cloud_api(body):
    req = pb.RemoveCloudAPIRequest()
    req.cloudName = body.get("cloudName", "")
    req.userName = body.get("userName", "")
    req.permanentRemove = body.get("permanentRemove", True)
    return _handle_call(_stub.RemoveCloudAPI, req)


def _get_space_info(body):
    req = pb.FileRequest()
    req.path = body.get("cloudName", "")
    return _handle_call(_stub.GetSpaceInfo, req)


def _get_cloud_api_config(body):
    req = pb.GetCloudAPIConfigRequest()
    req.cloudName = body.get("cloudName", "")
    return _handle_call(_stub.GetCloudAPIConfig, req)


def _login_webdav(body):
    _ensure_connected()
    try:
        req = pb.LoginWebDavRequest()
        req.serverUrl = body["serverUrl"]
        req.userName = body["userName"]
        req.password = body["password"]
        req.doNotSyncToCloud = body.get("doNotSyncToCloud", False)
        return _handle_call(_stub.APILoginWebDav, req)
    except Exception as e:
        return {"success": False, "error": str(e)}


def _get_account_status(body):
    return _handle_call(_stub.GetAccountStatus, Empty())


def _get_online_devices(body):
    return _handle_call(_stub.GetOnlineDevices, Empty())


def _create_folder(body):
    _ensure_connected()
    metadata = _auth_metadata()
    try:
        req = pb.CreateFolderRequest()
        req.parentPath = body.get("parentPath", "")
        req.folderName = body.get("folderName", "")
        resp = _stub.CreateFolder(req, timeout=10, metadata=metadata)
        return {"success": True, "data": _proto_to_dict(resp)}
    except Exception as e:
        return {"success": False, "error": str(e)}


def _delete_file(body):
    return _handle_call(_stub.DeleteFile,
                        pb.FileRequest(path=body.get("path", "")))


def _rename_file(body):
    _ensure_connected()
    metadata = _auth_metadata()
    try:
        req = pb.RenameFileRequest()
        req.theFilePath = body.get("path", "")
        req.newName = body.get("newName", "")
        resp = _stub.RenameFile(req, timeout=10, metadata=metadata)
        return {"success": True, "data": _proto_to_dict(resp)}
    except Exception as e:
        return {"success": False, "error": str(e)}


def _discover_smb_servers(body):
    _ensure_connected()
    try:
        resp = _stub.DiscoverSmbServers(Empty(), timeout=15, metadata=_auth_metadata())
        return {"success": True, "data": _proto_to_dict(resp)}
    except Exception as e:
        return {"success": False, "error": str(e)}


def _discover_smb_shares(body):
    _ensure_connected()
    try:
        req = pb.DiscoverSmbSharesRequest()
        req.server = body["server"]
        req.port = body.get("port", 445)
        req.userName = body["userName"]
        req.password = body["password"]
        if "workgroup" in body:
            req.workgroup = body["workgroup"]
        resp = _stub.DiscoverSmbShares(req, timeout=15, metadata=_auth_metadata())
        return {"success": True, "data": _proto_to_dict(resp)}
    except Exception as e:
        return {"success": False, "error": str(e)}


def _login_smb(body):
    _ensure_connected()
    try:
        req = pb.LoginSmbRequest()
        req.server = body["server"]
        req.share = body["share"]
        req.port = body.get("port", 445)
        req.userName = body["userName"]
        req.password = body["password"]
        req.doNotSyncToCloud = body.get("doNotSyncToCloud", False)
        if "workgroup" in body:
            req.workgroup = body["workgroup"]
        if "rootPath" in body:
            req.rootPath = body["rootPath"]
        resp = _stub.APILoginSmb(req, timeout=10, metadata=_auth_metadata())
        return {"success": True, "data": _proto_to_dict(resp)}
    except Exception as e:
        return {"success": False, "error": str(e)}


def _search_files(body):
    """Collect all results from the GetSearchResults stream."""
    _ensure_connected()
    try:
        req = pb.SearchRequest()
        req.path = body.get("path", "/")
        req.searchFor = body.get("searchFor", "")
        req.forceRefresh = body.get("forceRefresh", False)
        req.fuzzyMatch = body.get("fuzzyMatch", True)
        if "contentSearch" in body:
            req.contentSearch = body["contentSearch"]
        stream = _stub.GetSearchResults(req, timeout=30, metadata=_auth_metadata())
        files = []
        for resp in stream:
            if resp.subFiles:
                for f in resp.subFiles:
                    files.append(_proto_to_dict(f))
        return {"success": True, "data": files}
    except Exception as e:
        return {"success": False, "error": str(e)}


# ---- QR Login Handlers (session-based) ----

def _qr_login_start(body):
    """Start a QR login stream. Returns first message immediately with a session ID."""
    provider = body.get("provider", "")
    timeout = body.get("timeout", 120)
    _ensure_connected()
    metadata = _auth_metadata()

    try:
        if provider == "115open":
            req = pb.Login115OpenQRCodeRequest()
            stream = _stub.APILogin115OpenQRCode(req, timeout=timeout, metadata=metadata)
        elif provider == "aliyundrive":
            req = pb.LoginAliyundriveQRCodeRequest()
            req.useOpenAPI = True
            stream = _stub.APILoginAliyunDriveQRCode(req, timeout=timeout, metadata=metadata)
        elif provider == "189":
            req = pb.Login189QRCodeRequest()
            stream = _stub.APILogin189QRCode(req, timeout=timeout, metadata=metadata)
        elif provider == "guangyapan":
            req = pb.LoginGuangYaPanQRCodeRequest()
            stream = _stub.APILoginGuangYaPanQRCode(req, timeout=timeout, metadata=metadata)
        else:
            return {"success": False, "error": f"Unknown provider: {provider}"}
    except Exception as e:
        return {"success": False, "error": str(e)}

    session_id = uuid.uuid4().hex[:8]
    q = queue.Queue()

    def _stream_worker():
        try:
            for msg in stream:
                q.put({"__msg__": _proto_to_dict(msg)})
        except Exception as e:
            q.put({"__error__": str(e)})
        finally:
            q.put({"__done__": True})

    t = threading.Thread(target=_stream_worker, daemon=True)
    t.start()
    _qr_sessions[session_id] = {"queue": q, "thread": t, "stream": stream}

    # Wait for the first message (QR code) — up to 15 seconds
    try:
        first = q.get(timeout=15)
    except queue.Empty:
        _qr_sessions.pop(session_id, None)
        return {"success": False, "error": "Timeout waiting for QR code from server"}

    if first.get("__error__"):
        _qr_sessions.pop(session_id, None)
        return {"success": False, "error": first["__error__"]}
    if first.get("__done__"):
        _qr_sessions.pop(session_id, None)
        return {"success": False, "error": "Stream ended before QR code was received"}

    msg = first["__msg__"]
    msg["session"] = session_id
    return {"success": True, "data": msg}


def _qr_login_poll(body):
    """Poll a QR login session for the next message.

    Returns data with a `closed` flag embedded so the client can check it.
    data.closed = True  → session ended (success or error), stop polling
    data.closed = False, data.messageType present → new QR message
    data.closed = False, no messageType → no update yet, keep polling
    """
    session_id = body.get("session", "")
    session = _qr_sessions.get(session_id)
    if not session:
        return {"success": True, "data": {"closed": True, "error": "Session not found"}}

    try:
        item = session["queue"].get(timeout=0.5)
    except queue.Empty:
        return {"success": True, "data": {"closed": False}}

    if item.get("__error__"):
        _qr_sessions.pop(session_id, None)
        return {"success": True, "data": {"closed": True, "error": item["__error__"]}}
    if item.get("__done__"):
        _qr_sessions.pop(session_id, None)
        return {"success": True, "data": {"closed": True}}

    msg = item["__msg__"]
    msg["closed"] = False
    return {"success": True, "data": msg}


def _qr_login_cancel(body):
    """Cancel a QR login session."""
    session_id = body.get("session", "")
    session = _qr_sessions.pop(session_id, None)
    if session:
        # Cancel the gRPC stream to stop the background thread
        try:
            session.get("stream").cancel()
        except Exception:
            pass
    return {"success": True}


def _clear_token():
    """Clear the in-memory token."""
    global _token
    _token = None
    return {"success": True}


# ---- Route table ----

ROUTES = {
    "/SystemInfo":           _sys_info,
    "/Register":             _register,
    "/Login":                _login,
    "/GetToken":             _get_token,
    "/Logout":               _logout,
    "/RuntimeInfo":          _runtime_info,
    "/RunningInfo":          _running_info,
    "/SystemSettings":       _system_settings,
    "/SetSystemSettings":    _set_system_settings,
    "/GetAllCloudApis":      _get_all_cloud_apis,
    "/GetCloudMemberships":  _get_cloud_memberships,
    "/CanAddMoreCloudApis":  _can_add_more,
    "/GetSubFiles":          _get_sub_files,
    "/GetDownloadUrl":       _get_download_url,
    "/GetMountPoints":       _get_mount_points,
    "/AddMountPoint":        _add_mount_point,
    "/RemoveMountPoint":     _remove_mount_point,
    "/Mount":                _mount,
    "/Unmount":              _unmount,
    "/RemoveCloudAPI":       _remove_cloud_api,
    "/GetSpaceInfo":         _get_space_info,
    "/GetCloudAPIConfig":    _get_cloud_api_config,
    "/GetAccountStatus":     _get_account_status,
    "/GetOnlineDevices":     _get_online_devices,
    "/CreateFolder":         _create_folder,
    "/DeleteFile":           _delete_file,
    "/RenameFile":           _rename_file,
    "/SearchFiles":          _search_files,
    "/LoginWebDav":          _login_webdav,
    "/DiscoverSmbServers":   _discover_smb_servers,
    "/DiscoverSmbShares":    _discover_smb_shares,
    "/LoginSmb":             _login_smb,
    "/QRLoginStart":         _qr_login_start,
    "/QRLoginPoll":          _qr_login_poll,
    "/QRLoginCancel":        _qr_login_cancel,
    "/ClearToken":          lambda b: _clear_token(),
    "/HasUpdate":            lambda b: _handle_call(_stub.HasUpdate, Empty()),
    "/CheckUpdate":          lambda b: _handle_call(_stub.CheckUpdate, Empty()),
    "/DownloadUpdate":       lambda b: _handle_call(_stub.DownloadUpdate, Empty()),
    "/UpdateSystem":         lambda b: _handle_call(_stub.UpdateSystem, Empty()),
}


class ProxyHandler(BaseHTTPRequestHandler):

    def do_POST(self):
        # Read request body
        body = {}
        content_length = int(self.headers.get("Content-Length", 0))
        if content_length > 0:
            raw = self.rfile.read(content_length)
            try:
                body = json.loads(raw)
            except json.JSONDecodeError:
                self._json_reply(400, {"success": False, "error": "Invalid JSON"})
                return

        handler = ROUTES.get(self.path)
        if handler is None:
            self._json_reply(404, {"success": False,
                                   "error": f"Unknown endpoint: {self.path}"})
            return

        try:
            result = handler(body)
            self._json_reply(200, result)
        except Exception as e:
            traceback.print_exc()
            self._json_reply(500, {"success": False, "error": str(e)})

    def do_GET(self):
        if self.path == "/health":
            self._json_reply(200, {"status": "ok"})
        elif self.path == "/shutdown":
            self._json_reply(200, {"status": "shutting_down"})
            # Shutdown in a separate thread to allow response to be sent
            import threading
            threading.Thread(target=self.server.shutdown).start()
        else:
            self._json_reply(404, {"success": False, "error": "Not found"})

    def _json_reply(self, code, data):
        body = json.dumps(data, ensure_ascii=False).encode("utf-8")
        self.send_response(code)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def log_message(self, format, *args):
        """Suppress default HTTP request logging to stderr."""
        pass


def main():
    print(f"[CD2 Proxy] Starting on {HOST}:{PORT}")
    print(f"[CD2 Proxy] gRPC target: {CD2_HOST}:{CD2_PORT}")
    server = HTTPServer((HOST, PORT), ProxyHandler)
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        pass
    finally:
        if _channel:
            _channel.close()
        server.server_close()
        print("[CD2 Proxy] Stopped")


if __name__ == "__main__":
    main()
