from sari.mcp.cli.smart_daemon import ensure_smart_daemon
from sari.mcp.trace import trace
from sari.core.daemon_resolver import resolve_daemon_address as _resolve_daemon_target
from sari.core.workspace import WorkspaceManager
from sari.mcp.telemetry import TelemetryLogger
import sys
import json
import socket
import threading
import os
import time
import logging
import secrets
from pathlib import Path

# Add project root to sys.path for absolute imports
SCRIPT_DIR = Path(__file__).parent
REPO_ROOT = SCRIPT_DIR.parent
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))


try:
    import fcntl  # type: ignore
except Exception:
    fcntl = None

# Configure logging to stderr so it doesn't interfere with MCP STDIO
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    stream=sys.stderr
)
logger = logging.getLogger("mcp-proxy")
telemetry = TelemetryLogger(WorkspaceManager.get_global_log_dir())


def _log_info(message: str) -> None:
    logger.info(message)
    try:
        telemetry.log_info(message)
    except Exception:
        pass


def _log_error(message: str) -> None:
    logger.error(message)
    try:
        telemetry.log_error(message)
    except Exception:
        pass


DEFAULT_HOST = "127.0.0.1"
DEFAULT_PORT = 47779
MAX_MESSAGE_SIZE = 10 * 1024 * 1024  # 10MB
_HEADER_SEP = b"\r\n\r\n"
_MODE_FRAMED = "framed"
_MODE_JSONL = "jsonl"


def _identify_sari_daemon(host: str, port: int, timeout: float = 0.3) -> bool:
    try:
        trace("proxy_identify_start", host=host, port=port)
        with socket.create_connection((host, port), timeout=timeout) as sock:
            sock.settimeout(timeout)
            body = json.dumps({"jsonrpc": "2.0", "id": 1,
                              "method": "sari/identify"}).encode("utf-8")
            header = f"Content-Length: {len(body)}\r\n\r\n".encode("ascii")
            sock.sendall(header + body)

            f = sock.makefile("rb")
            headers = {}
            while True:
                line = f.readline()
                if not line:
                    return False
                line = line.strip()
                if not line:
                    break
                if b":" in line:
                    k, v = line.split(b":", 1)
                    headers[k.strip().lower()] = v.strip()

            try:
                content_length = int(headers.get(b"content-length", b"0"))
            except ValueError:
                return False
            if content_length <= 0:
                return False
            resp_body = f.read(content_length)
            if not resp_body:
                return False
            resp = json.loads(resp_body.decode("utf-8"))
            result = resp.get("result") or {}
            if result.get("name") == "sari":
                trace("proxy_identify_ok", host=host, port=port)
                return True
    except Exception:
        trace("proxy_identify_error", host=host, port=port)
        pass
    return False


def _lock_file(lock_file) -> None:
    if fcntl is not None:
        fcntl.flock(lock_file, fcntl.LOCK_EX)
        return
    try:
        import msvcrt
        lock_file.seek(0)
        msvcrt.locking(lock_file.fileno(), msvcrt.LK_LOCK, 1)
    except Exception:
        pass


def _unlock_file(lock_file) -> None:
    if fcntl is not None:
        fcntl.flock(lock_file, fcntl.LOCK_UN)
        return
    try:
        import msvcrt
        lock_file.seek(0)
        msvcrt.locking(lock_file.fileno(), msvcrt.LK_UNLCK, 1)
    except Exception:
        pass


def start_daemon_if_needed(host, port, workspace_root: str = ""):
    """Ensures daemon is running using smart logic."""
    try:
        ensure_smart_daemon(host, port, workspace_root)
        return True
    except Exception as e:
        _log_error(f"Failed to start daemon at {host}:{port}: {e}")
        trace("proxy_start_daemon_error", host=host, port=port, error=str(e))
        return False


def forward_socket_to_stdout(sock, state):
    try:
        f = sock.makefile("rb")
        while True:
            # Read Headers
            headers = {}
            while True:
                line = f.readline()
                if not line:
                    break
                line_str = line.decode("utf-8").strip()
                if not line_str:
                    break
                if ":" in line_str:
                    k, v = line_str.split(":", 1)
                    headers[k.strip().lower()] = v.strip()

            if not headers and not line:
                break

            content_length = int(headers.get("content-length", 0))
            if content_length <= 0:
                continue

            body = f.read(content_length)
            if not body:
                break

            msg_id = None
            obj = None
            try:
                obj = json.loads(body.decode("utf-8"))
                if isinstance(obj, dict):
                    msg_id = obj.get("id")
            except Exception:
                msg_id = None
                obj = None

            # If daemon is draining, reconnect and replay initialize without
            # surfacing an error to the client.
            try:
                if isinstance(obj, dict):
                    err = obj.get("error") or {}
                    code = err.get("code")
                    message = str(err.get("message") or "")
                    if code == -32001 and "draining" in message.lower():
                        _log_info(
                            "Server draining detected; reconnecting to latest daemon.")
                        state["dead"] = True
                        if _reconnect(state):
                            init_req = state.get("init_request")
                            if init_req:
                                _send_payload(
                                    state,
                                    json.dumps(init_req).encode("utf-8"),
                                    state.get("mode") or _MODE_FRAMED,
                                )
                            continue
            except Exception:
                pass
            if msg_id is not None:
                with state["suppress_lock"]:
                    suppress_ids = state.setdefault("suppress_ids", set())
                    if msg_id in suppress_ids:
                        suppress_ids.discard(msg_id)
                        continue

            mode = state.get("mode") or _MODE_FRAMED
            if mode == _MODE_JSONL:
                sys.stdout.buffer.write(body + b"\n")
                sys.stdout.buffer.flush()
            else:
                header = f"Content-Length: {len(body)}\r\n\r\n".encode("ascii")
                sys.stdout.buffer.write(header + body)
                sys.stdout.buffer.flush()
            trace("proxy_forwarded_to_stdout", bytes=len(body))
    except Exception as e:
        _log_error(f"Error forwarding socket to stdout: {e}")
        trace("proxy_forward_stdout_error", error=str(e))
    finally:
        state["dead"] = True
        try:
            sock.close()
        except Exception:
            pass


def _read_mcp_message(stdin):
    """Robustly reads one MCP message, skipping noise or leading logs."""
    while True:
        line = stdin.readline()
        if not line:
            return None

        stripped = line.strip()
        if not stripped:
            continue

        # 1. JSONL fallback (starts with '{')
        if stripped.startswith(b"{"):
            return stripped, _MODE_JSONL

        # 2. Content-Length framing
        if stripped.lower().startswith(b"content-length:"):
            try:
                # Re-parse headers fully
                content_length = None
                headers = [line]
                while True:
                    h = stdin.readline()
                    if not h or h in (b"\n", b"\r\n"):
                        break
                    headers.append(h)

                for h in headers:
                    if h.lower().startswith(b"content-length:"):
                        content_length = int(h.split(b":", 1)[1].strip())
                        break

                if content_length and 0 < content_length <= MAX_MESSAGE_SIZE:
                    body = b""
                    while len(body) < content_length:
                        chunk = stdin.read(content_length - len(body))
                        if not chunk:
                            break
                        body += chunk
                    if len(body) == content_length:
                        return body, _MODE_FRAMED
            except Exception:
                continue  # Malformed frame, keep looking

        # 3. If noise, keep looping
        continue


def _send_payload(state, payload: bytes, mode: str) -> None:
    sock = state.get("sock")
    if not sock:
        raise OSError("socket not connected")
    if mode == _MODE_JSONL:
        sock.sendall(payload + b"\n")
    else:
        header = f"Content-Length: {len(payload)}\r\n\r\n".encode("ascii")
        sock.sendall(header + payload)
    trace("proxy_send_payload", bytes=len(payload), mode=mode)


def _reconnect(state) -> bool:
    max_attempts = int(os.environ.get("SARI_PROXY_RECONNECT_MAX", "10") or 10)
    backoff = float(
        os.environ.get(
            "SARI_PROXY_RECONNECT_BACKOFF",
            "0.2") or 0.2)
    with state["conn_lock"]:
        if not state.get("dead") and state.get("sock"):
            return True

        workspace_root = state.get("workspace_root") or ""
        last_err = None
        for _ in range(max_attempts):
            host, port = _resolve_daemon_target()
            if not start_daemon_if_needed(
                    host, port, workspace_root=workspace_root):
                last_err = "daemon start failed"
                time.sleep(backoff)
                continue
            host, port = _resolve_daemon_target()
            try:
                sock = socket.create_connection((host, port))
                last_err = None
                break
            except Exception as e:
                last_err = str(e)
                time.sleep(backoff)
        else:
            _log_error(f"Reconnect failed: {last_err}")
            trace("proxy_reconnect_failed", error=last_err)
            return False

        state["sock"] = sock
        state["dead"] = False
        trace("proxy_reconnect_ok", host=host, port=port)

        t = threading.Thread(
            target=forward_socket_to_stdout, args=(
                sock, state), daemon=True)
        t.start()

        init_req = state.get("init_request")
        if init_req:
            try:
                internal = dict(init_req)
                with state["suppress_lock"]:
                    suppress_ids = state.setdefault("suppress_ids", set())
                    internal_id = -secrets.randbelow(2**31 - 1) - 1
                    while internal_id in suppress_ids:
                        internal_id = -secrets.randbelow(2**31 - 1) - 1
                    suppress_ids.add(internal_id)
                internal["id"] = internal_id
                _send_payload(state, json.dumps(internal).encode(
                    "utf-8"), state.get("mode") or _MODE_FRAMED)
            except Exception as e:
                _log_error(f"Reconnect initialize failed: {e}")
        return True


def forward_stdin_to_socket(state):
    try:
        stdin = sys.stdin.buffer
        while True:
            res = _read_mcp_message(stdin)
            if res is None:
                break
            msg, mode = res
            trace("proxy_received_from_stdin", bytes=len(msg), mode=mode)
            # Inject rootUri for initialize when client omits it
            # (per-connection workspace)
            try:
                req = json.loads(msg.decode("utf-8"))

                def _inject(obj):
                    if not isinstance(obj, dict) or obj.get(
                            "method") != "initialize":
                        return obj, False
                    params = obj.get("params") or {}
                    if params.get("rootUri") or params.get("rootPath"):
                        state["init_request"] = obj
                        return obj, False

                    # PRIORITY: SARI_
                    ws = os.environ.get("SARI_WORKSPACE_ROOT") or state.get(
                        "workspace_root")
                    if not ws:
                        state["init_request"] = obj
                        return obj, False
                    params = dict(params)
                    params["rootUri"] = f"file://{ws}"
                    obj = dict(obj)
                    obj["params"] = params
                    _log_info(f"Injected rootUri for initialize: {ws}")
                    state["init_request"] = obj
                    return obj, True

                injected = False
                if isinstance(req, dict):
                    req, injected = _inject(req)
                elif isinstance(req, list):
                    new_list = []
                    for item in req:
                        item2, did = _inject(item)
                        injected = injected or did
                        new_list.append(item2)
                    req = new_list

                if injected:
                    msg = json.dumps(req).encode("utf-8")
            except Exception:
                pass
            if state.get("mode") is None:
                state["mode"] = mode

            if state.get("dead"):
                if not _reconnect(state):
                    return
            try:
                with state["send_lock"]:
                    _send_payload(state, msg, mode)
            except Exception:
                state["dead"] = True
                if not _reconnect(state):
                    return
                with state["send_lock"]:
                    _send_payload(state, msg, mode)
    except Exception as e:
        _log_error(f"Error forwarding stdin to socket: {e}")
        trace("proxy_forward_stdin_error", error=str(e))
        try:
            sock = state.get("sock")
            if sock:
                sock.close()
        except Exception:
            pass


def _write_error_to_client(message: str, code: int = -32002):
    """Sends a formal JSON-RPC error directly to stdout."""
    resp = {
        "jsonrpc": "2.0",
        "error": {"code": code, "message": message}
    }
    payload = json.dumps(resp).encode("utf-8")
    header = f"Content-Length: {len(payload)}\r\n\r\n".encode("ascii")
    sys.stdout.buffer.write(header + payload)
    sys.stdout.buffer.flush()


def main():
    # Log startup context for diagnostics
    _log_info(
        "Proxy startup: cwd=%s argv=%s SARI_ROOT=%s"
        % (
            os.getcwd(),
            sys.argv,
            os.environ.get("SARI_WORKSPACE_ROOT"),
        )
    )
    trace("proxy_startup", cwd=os.getcwd(), argv=sys.argv,
          workspace_root=os.environ.get("SARI_WORKSPACE_ROOT"))

    workspace_root = os.environ.get(
        "SARI_WORKSPACE_ROOT") or WorkspaceManager.resolve_workspace_root()
    host, port = _resolve_daemon_target()

    if not start_daemon_if_needed(host, port, workspace_root=workspace_root):
        _write_error_to_client(
            "Sari Daemon failed to start. Run 'sari doctor' for diagnostics.")
        sys.exit(1)

    try:
        host, port = _resolve_daemon_target()
        sock = socket.create_connection((host, port))
    except Exception as e:
        err_msg = f"Could not connect to daemon at {host}:{port}: {e}"
        _log_error(err_msg)
        trace("proxy_connect_error", error=str(e))
        _write_error_to_client(err_msg)
        sys.exit(1)

    # Start threads for bidirectional forwarding
    state = {
        "mode": None,
        "sock": sock,
        "dead": False,
        "send_lock": threading.Lock(),
        "conn_lock": threading.Lock(),
        "suppress_lock": threading.Lock(),
        "suppress_ids": set(),
        "init_request": None,
        "workspace_root": workspace_root,
    }
    t1 = threading.Thread(
        target=forward_socket_to_stdout, args=(
            sock, state), daemon=True)
    t1.start()

    forward_stdin_to_socket(state)
    t1.join(timeout=1.0)


if __name__ == "__main__":
    main()
