import asyncio
import json
import shutil
import tempfile
from dataclasses import dataclass
from pathlib import Path

WORKSPACE_ROOT = Path(tempfile.gettempdir()) / "theplug-lsp"
CONFIG_FILE = "pyrightconfig.json"
ENTRY_FILE = "proxy_addon.py"

STUB_FILES = {
    "typings/mitmproxy/__init__.pyi": """from . import (
    command,
    ctx,
    dns,
    http,
    tcp,
    tls,
    udp,
    websocket,
)

__all__ = [
    "command",
    "ctx",
    "dns",
    "http",
    "tcp",
    "tls",
    "udp",
    "websocket",
]
""",
    "typings/mitmproxy/ctx.pyi": """class _Logger:
    def debug(self, message: object) -> None: ...
    def info(self, message: object) -> None: ...
    def warn(self, message: object) -> None: ...
    def error(self, message: object) -> None: ...
    def alert(self, message: object) -> None: ...


class _Options:
    def __getattr__(self, name: str) -> object: ...


log: _Logger
options: _Options
""",
    "typings/mitmproxy/http.pyi": """from typing import Any


class Headers(dict[str, str]):
    ...


class Request:
    host: str
    path: str
    headers: Headers
    query: dict[str, str]
    content: bytes


class Response:
    headers: Headers
    status_code: int
    content: bytes

    @classmethod
    def make(
        cls,
        status_code: int,
        content: bytes | str = b"",
        headers: dict[str, str] | None = None,
    ) -> "Response": ...


class WebSocketMessage:
    content: bytes
    from_client: bool


class WebSocketData:
    messages: list[WebSocketMessage]


class HTTPFlow:
    request: Request
    response: Response | None
    websocket: WebSocketData | None
    metadata: dict[str, Any]
""",
    "typings/mitmproxy/tcp.pyi": """class TCPMessage:
    content: bytes
    from_client: bool


class TCPFlow:
    messages: list[TCPMessage]
""",
    "typings/mitmproxy/dns.pyi": """class Question:
    name: str
    type: int


class ResourceRecord:
    def __init__(
        self,
        *,
        name: str,
        type: int,
        class_: int,
        ttl: int,
        data: bytes,
    ) -> None: ...


class Message:
    question: Question | None

    def succeed(self, records: list[ResourceRecord]) -> "Message": ...


class _Types:
    A: int


class _Classes:
    IN: int


types: _Types
classes: _Classes


class DNSFlow:
    request: Message
    response: Message | None
""",
    "typings/mitmproxy/command.pyi": """from collections.abc import Callable
from typing import TypeVar

F = TypeVar("F", bound=Callable[..., object])


def command(name: str) -> Callable[[F], F]: ...
""",
    "typings/mitmproxy/websocket.pyi": """class WebSocketMessage:
    content: bytes
    from_client: bool
""",
    "typings/mitmproxy/tls.pyi": """class ClientHelloData: ...
""",
    "typings/mitmproxy/udp.pyi": """class UDPMessage:
    content: bytes
    from_client: bool
""",
}


class LspFrameBuffer:
    def __init__(self) -> None:
        self.buffer = bytearray()

    def feed(self, chunk: bytes) -> list[str]:
        self.buffer.extend(chunk)
        messages: list[str] = []
        while True:
            header_end = self.buffer.find(b"\r\n\r\n")
            if header_end < 0:
                return messages

            header = self.buffer[:header_end].decode("ascii", errors="ignore")
            size = content_length(header)
            if size is None:
                del self.buffer[: header_end + 4]
                continue

            frame_end = header_end + 4 + size
            if len(self.buffer) < frame_end:
                return messages

            body = bytes(self.buffer[header_end + 4 : frame_end])
            del self.buffer[:frame_end]
            messages.append(body.decode("utf-8"))


@dataclass(slots=True)
class PythonEditorSession:
    root_path: Path
    file_path: Path

    @property
    def root_uri(self) -> str:
        return self.root_path.as_uri()

    @property
    def file_uri(self) -> str:
        return self.file_path.as_uri()


def content_length(header: str) -> int | None:
    for line in header.split("\r\n"):
        if ":" not in line:
            continue
        name, value = line.split(":", 1)
        if name.lower() == "content-length":
            return int(value.strip())
    return None


def encode_message(message: str) -> bytes:
    payload = message.encode("utf-8")
    header = f"Content-Length: {len(payload)}\r\n\r\n".encode("ascii")
    return header + payload


def lsp_binary() -> str | None:
    return shutil.which("basedpyright-langserver")


def get_session() -> PythonEditorSession:
    root_path = WORKSPACE_ROOT
    file_path = root_path / ENTRY_FILE

    root_path.mkdir(parents=True, exist_ok=True)
    file_path.touch(exist_ok=True)

    config = {
        "include": [ENTRY_FILE],
        "stubPath": "typings",
        "typeCheckingMode": "basic",
        "pythonVersion": "3.14",
        "reportMissingTypeStubs": False,
        "reportMissingImports": "warning",
        "useLibraryCodeForTypes": True,
    }
    (root_path / CONFIG_FILE).write_text(
        json.dumps(config, indent=2),
        encoding="utf-8",
    )

    for relative_path, content in STUB_FILES.items():
        path = root_path / relative_path
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_text(content, encoding="utf-8")

    return PythonEditorSession(root_path=root_path, file_path=file_path)


async def open_process(session: PythonEditorSession) -> asyncio.subprocess.Process:
    binary = lsp_binary()
    if binary is None:
        raise RuntimeError("basedpyright-langserver is not installed")

    return await asyncio.create_subprocess_exec(
        binary,
        "--stdio",
        cwd=session.root_path,
        stdin=asyncio.subprocess.PIPE,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
    )


async def forward_client(websocket, writer: asyncio.StreamWriter) -> None:
    while True:
        message = await websocket.receive_text()
        writer.write(encode_message(message))
        await writer.drain()


async def forward_server(reader: asyncio.StreamReader, websocket) -> None:
    frames = LspFrameBuffer()
    while True:
        chunk = await reader.read(65536)
        if not chunk:
            return
        for message in frames.feed(chunk):
            await websocket.send_text(message)


async def drain_errors(reader: asyncio.StreamReader) -> None:
    while True:
        chunk = await reader.read(4096)
        if not chunk:
            return


async def bridge(websocket, session: PythonEditorSession) -> None:
    process = await open_process(session)
    stdin = process.stdin
    stdout = process.stdout
    stderr = process.stderr

    if stdin is None or stdout is None or stderr is None:
        process.kill()
        await process.wait()
        raise RuntimeError("basedpyright-langserver did not expose stdio pipes")

    send_task = asyncio.create_task(forward_client(websocket, stdin))
    receive_task = asyncio.create_task(forward_server(stdout, websocket))
    error_task = asyncio.create_task(drain_errors(stderr))

    try:
        done, pending = await asyncio.wait(
            {send_task, receive_task, error_task},
            return_when=asyncio.FIRST_COMPLETED,
        )
        for task in pending:
            task.cancel()
        for task in done:
            task.result()
    finally:
        if process.returncode is None:
            stdin.close()
            await stdin.wait_closed()
            try:
                await asyncio.wait_for(process.wait(), timeout=1)
            except TimeoutError:
                process.kill()
                await process.wait()