import hashlib
import json
import secrets
import threading
import time
import uuid
import logging
from dataclasses import dataclass
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from typing import Any, Callable, Optional

from .stata_client import StataClient
from .config import (
    DEFAULT_HOST,
    DEFAULT_PORT,
    MAX_ARROW_LIMIT,
    MAX_CHARS,
    MAX_LIMIT,
    MAX_REQUEST_BYTES,
    MAX_VARS,
    TOKEN_TTL_S,
    VIEW_TTL_S,
)


logger = logging.getLogger("mcp_stata")


def _stable_hash(payload: dict[str, Any]) -> str:
    return hashlib.sha1(json.dumps(payload, sort_keys=True).encode("utf-8")).hexdigest()


@dataclass
class UIChannelInfo:
    base_url: str
    token: str
    expires_at: int


@dataclass
class ViewHandle:
    view_id: str
    dataset_id: str
    frame: str
    filter_expr: str
    obs_indices: list[int]
    filtered_n: int
    created_at: float
    last_access: float


class UIChannelManager:
    def __init__(
        self,
        client: StataClient,
        *,
        host: str = DEFAULT_HOST,
        port: int = DEFAULT_PORT,
        token_ttl_s: int = TOKEN_TTL_S,
        view_ttl_s: int = VIEW_TTL_S,
        max_limit: int = MAX_LIMIT,
        max_vars: int = MAX_VARS,
        max_chars: int = MAX_CHARS,
        max_request_bytes: int = MAX_REQUEST_BYTES,
        max_arrow_limit: int = MAX_ARROW_LIMIT,
    ):
        self._client = client
        self._host = host
        self._port = port
        self._token_ttl_s = token_ttl_s
        self._view_ttl_s = view_ttl_s
        self._max_limit = max_limit
        self._max_vars = max_vars
        self._max_chars = max_chars
        self._max_request_bytes = max_request_bytes
        self._max_arrow_limit = max_arrow_limit

        self._lock = threading.Lock()
        self._httpd: ThreadingHTTPServer | None = None
        self._thread: threading.Thread | None = None

        self._token: str | None = None
        self._expires_at: int = 0

        self._dataset_version: int = 0
        self._dataset_id_cache: str | None = None
        self._dataset_id_cache_at_version: int = -1

        self._views: dict[str, ViewHandle] = {}

    def notify_potential_dataset_change(self) -> None:
        with self._lock:
            self._dataset_version += 1
            self._dataset_id_cache = None
            self._views.clear()

    def get_channel(self) -> UIChannelInfo:
        self._ensure_http_server()
        with self._lock:
            self._ensure_token()
            assert self._httpd is not None
            port = self._httpd.server_address[1]
            base_url = f"http://{self._host}:{port}"
            return UIChannelInfo(base_url=base_url, token=self._token or "", expires_at=self._expires_at)

    def capabilities(self) -> dict[str, bool]:
        return {"dataBrowser": True, "filtering": True, "sorting": True, "arrowStream": True}

    def current_dataset_id(self) -> str:
        with self._lock:
            if self._dataset_id_cache is not None and self._dataset_id_cache_at_version == self._dataset_version:
                return self._dataset_id_cache

        state = self._client.get_dataset_state()
        payload = {
            "version": self._dataset_version,
            "frame": state.get("frame"),
            "n": state.get("n"),
            "k": state.get("k"),
            "sortlist": state.get("sortlist"),
        }
        digest = _stable_hash(payload)

        with self._lock:
            self._dataset_id_cache = digest
            self._dataset_id_cache_at_version = self._dataset_version
            return digest

    def get_view(self, view_id: str) -> Optional[ViewHandle]:
        now = time.time()
        with self._lock:
            self._evict_expired_locked(now)
            view = self._views.get(view_id)
            if view is None:
                return None
            view.last_access = now
            return view

    def create_view(self, *, dataset_id: str, frame: str, filter_expr: str) -> ViewHandle:
        current_id = self.current_dataset_id()
        if dataset_id != current_id:
            raise DatasetChangedError(current_id)

        try:
            obs_indices = self._client.compute_view_indices(filter_expr)
        except ValueError as e:
            raise InvalidFilterError(str(e))
        except RuntimeError as e:
            msg = str(e) or "No data in memory"
            if "no data" in msg.lower():
                raise NoDataInMemoryError(msg)
            raise
        now = time.time()
        view_id = f"view_{uuid.uuid4().hex}"
        view = ViewHandle(
            view_id=view_id,
            dataset_id=current_id,
            frame=frame,
            filter_expr=filter_expr,
            obs_indices=obs_indices,
            filtered_n=len(obs_indices),
            created_at=now,
            last_access=now,
        )
        with self._lock:
            self._evict_expired_locked(now)
            self._views[view_id] = view
        return view

    def delete_view(self, view_id: str) -> bool:
        with self._lock:
            return self._views.pop(view_id, None) is not None

    def validate_token(self, header_value: str | None) -> bool:
        if not header_value:
            return False
        if not header_value.startswith("Bearer "):
            return False
        token = header_value[len("Bearer ") :].strip()
        with self._lock:
            self._ensure_token()
            if self._token is None:
                return False
            if time.time() * 1000 >= self._expires_at:
                return False
            return secrets.compare_digest(token, self._token)

    def limits(self) -> tuple[int, int, int, int]:
        return self._max_limit, self._max_vars, self._max_chars, self._max_request_bytes

    def _ensure_token(self) -> None:
        now_ms = int(time.time() * 1000)
        if self._token is None or now_ms >= self._expires_at:
            self._token = secrets.token_urlsafe(32)
            self._expires_at = int((time.time() + self._token_ttl_s) * 1000)

    def _evict_expired_locked(self, now: float) -> None:
        expired: list[str] = []
        for key, view in self._views.items():
            if now - view.last_access >= self._view_ttl_s:
                expired.append(key)
        for key in expired:
            self._views.pop(key, None)

    def _ensure_http_server(self) -> None:
        with self._lock:
            if self._httpd is not None:
                return

            manager = self

            class Handler(BaseHTTPRequestHandler):

                def _send_json(self, status: int, payload: dict[str, Any]) -> None:
                    data = json.dumps(payload).encode("utf-8")
                    self.send_response(status)
                    self.send_header("Content-Type", "application/json")
                    self.send_header("Content-Length", str(len(data)))
                    self.end_headers()
                    self.wfile.write(data)

                def _send_binary(self, status: int, data: bytes, content_type: str) -> None:
                    self.send_response(status)
                    self.send_header("Content-Type", content_type)
                    self.send_header("Content-Length", str(len(data)))
                    self.end_headers()
                    self.wfile.write(data)

                def _error(self, status: int, code: str, message: str, *, stata_rc: int | None = None) -> None:
                    if status >= 500 or code == "internal_error":
                        logger.error("UI HTTP error %s: %s", code, message)
                        message = "Internal server error"
                    body: dict[str, Any] = {"error": {"code": code, "message": message}}
                    if stata_rc is not None:
                        body["error"]["stataRc"] = stata_rc
                    self._send_json(status, body)

                def _require_auth(self) -> bool:
                    if manager.validate_token(self.headers.get("Authorization")):
                        return True
                    self._error(401, "auth_failed", "Unauthorized")
                    return False

                def _read_json(self) -> dict[str, Any] | None:
                    max_limit, max_vars, max_chars, max_bytes = manager.limits()
                    _ = (max_limit, max_vars, max_chars)

                    length = int(self.headers.get("Content-Length", "0") or "0")
                    if length <= 0:
                        return {}
                    if length > max_bytes:
                        self._error(400, "request_too_large", "Request too large")
                        return None
                    raw = self.rfile.read(length)
                    try:
                        parsed = json.loads(raw.decode("utf-8"))
                    except Exception:
                        self._error(400, "invalid_request", "Invalid JSON")
                        return None
                    if not isinstance(parsed, dict):
                        self._error(400, "invalid_request", "Expected JSON object")
                        return None
                    return parsed

                def do_GET(self) -> None:
                    if not self._require_auth():
                        return

                    if self.path == "/v1/dataset":
                        try:
                            state = manager._client.get_dataset_state()
                            dataset_id = manager.current_dataset_id()
                            self._send_json(
                                200,
                                {
                                    "dataset": {
                                        "id": dataset_id,
                                        "frame": state.get("frame"),
                                        "n": state.get("n"),
                                        "k": state.get("k"),
                                        "changed": state.get("changed"),
                                    }
                                },
                            )
                            return
                        except NoDataInMemoryError as e:
                            self._error(400, "no_data_in_memory", str(e), stata_rc=e.stata_rc)
                            return
                        except Exception as e:
                            self._error(500, "internal_error", str(e))
                            return

                    if self.path == "/v1/vars":
                        try:
                            state = manager._client.get_dataset_state()
                            dataset_id = manager.current_dataset_id()
                            variables = manager._client.list_variables_rich()
                            self._send_json(
                                200,
                                {
                                    "dataset": {"id": dataset_id, "frame": state.get("frame")},
                                    "variables": variables,
                                },
                            )
                            return
                        except NoDataInMemoryError as e:
                            self._error(400, "no_data_in_memory", str(e), stata_rc=e.stata_rc)
                            return
                        except Exception as e:
                            self._error(500, "internal_error", str(e))
                            return

                    self._error(404, "not_found", "Not found")

                def do_POST(self) -> None:
                    if not self._require_auth():
                        return


                    if self.path == "/v1/arrow":
                        body = self._read_json()
                        if body is None:
                            return
                        try:
                            resp_bytes = handle_arrow_request(manager, body, view_id=None)
                            self._send_binary(200, resp_bytes, "application/vnd.apache.arrow.stream")
                            return
                        except HTTPError as e:
                            self._error(e.status, e.code, e.message, stata_rc=e.stata_rc)
                            return
                        except Exception as e:
                            self._error(500, "internal_error", str(e))
                            return

                    if self.path == "/v1/page":
                        body = self._read_json()
                        if body is None:
                            return
                        try:
                            resp = handle_page_request(manager, body, view_id=None)
                            self._send_json(200, resp)
                            return
                        except HTTPError as e:
                            self._error(e.status, e.code, e.message, stata_rc=e.stata_rc)
                            return
                        except Exception as e:
                            self._error(500, "internal_error", str(e))
                            return

                    if self.path == "/v1/views":
                        body = self._read_json()
                        if body is None:
                            return
                        dataset_id = str(body.get("datasetId", ""))
                        frame = str(body.get("frame", "default"))
                        filter_expr = str(body.get("filterExpr", ""))
                        if not dataset_id or not filter_expr:
                            self._error(400, "invalid_request", "datasetId and filterExpr are required")
                            return
                        try:
                            view = manager.create_view(dataset_id=dataset_id, frame=frame, filter_expr=filter_expr)
                            self._send_json(
                                200,
                                {
                                    "dataset": {"id": view.dataset_id, "frame": view.frame},
                                    "view": {"id": view.view_id, "filteredN": view.filtered_n},
                                },
                            )
                            return
                        except DatasetChangedError as e:
                            self._error(409, "dataset_changed", "Dataset changed")
                            return
                        except ValueError as e:
                            self._error(400, "invalid_filter", str(e))
                            return
                        except RuntimeError as e:
                            msg = str(e) or "No data in memory"
                            if "no data" in msg.lower():
                                self._error(400, "no_data_in_memory", msg)
                                return
                            self._error(500, "internal_error", msg)
                            return
                        except Exception as e:
                            self._error(500, "internal_error", str(e))
                            return

                    if self.path.startswith("/v1/views/") and self.path.endswith("/page"):
                        parts = self.path.split("/")
                        if len(parts) != 5:
                            self._error(404, "not_found", "Not found")
                            return
                        view_id = parts[3]
                        body = self._read_json()
                        if body is None:
                            return
                        try:
                            resp = handle_page_request(manager, body, view_id=view_id)
                            self._send_json(200, resp)
                            return
                        except HTTPError as e:
                            self._error(e.status, e.code, e.message, stata_rc=e.stata_rc)
                            return
                        except Exception as e:
                            self._error(500, "internal_error", str(e))
                            return

                    if self.path.startswith("/v1/views/") and self.path.endswith("/arrow"):
                        parts = self.path.split("/")
                        if len(parts) != 5:
                            self._error(404, "not_found", "Not found")
                            return
                        view_id = parts[3]
                        body = self._read_json()
                        if body is None:
                            return
                        try:
                            resp_bytes = handle_arrow_request(manager, body, view_id=view_id)
                            self._send_binary(200, resp_bytes, "application/vnd.apache.arrow.stream")
                            return
                        except HTTPError as e:
                            self._error(e.status, e.code, e.message, stata_rc=e.stata_rc)
                            return
                        except Exception as e:
                            self._error(500, "internal_error", str(e))
                            return

                    if self.path == "/v1/filters/validate":
                        body = self._read_json()
                        if body is None:
                            return
                        filter_expr = str(body.get("filterExpr", ""))
                        if not filter_expr:
                            self._error(400, "invalid_request", "filterExpr is required")
                            return
                        try:
                            manager._client.validate_filter_expr(filter_expr)
                            self._send_json(200, {"ok": True})
                            return
                        except ValueError as e:
                            self._error(400, "invalid_filter", str(e))
                            return
                        except RuntimeError as e:
                            msg = str(e) or "No data in memory"
                            if "no data" in msg.lower():
                                self._error(400, "no_data_in_memory", msg)
                                return
                            self._error(500, "internal_error", msg)
                            return
                        except Exception as e:
                            self._error(500, "internal_error", str(e))
                            return

                    self._error(404, "not_found", "Not found")

                def do_DELETE(self) -> None:
                    if not self._require_auth():
                        return

                    if self.path.startswith("/v1/views/"):
                        parts = self.path.split("/")
                        if len(parts) != 4:
                            self._error(404, "not_found", "Not found")
                            return
                        view_id = parts[3]
                        if manager.delete_view(view_id):
                            self._send_json(200, {"ok": True})
                        else:
                            self._error(404, "not_found", "Not found")
                        return

                    self._error(404, "not_found", "Not found")

                def log_message(self, format: str, *args: Any) -> None:
                    return

            httpd = ThreadingHTTPServer((self._host, self._port), Handler)
            t = threading.Thread(target=httpd.serve_forever, daemon=True)
            t.start()
            self._httpd = httpd
            self._thread = t


class HTTPError(Exception):
    def __init__(self, status: int, code: str, message: str, *, stata_rc: int | None = None):
        super().__init__(message)
        self.status = status
        self.code = code
        self.message = message
        self.stata_rc = stata_rc


class DatasetChangedError(Exception):
    def __init__(self, current_dataset_id: str):
        super().__init__("dataset_changed")
        self.current_dataset_id = current_dataset_id


class NoDataInMemoryError(Exception):
    def __init__(self, message: str = "No data in memory", *, stata_rc: int | None = None):
        super().__init__(message)
        self.stata_rc = stata_rc


class InvalidFilterError(Exception):
    def __init__(self, message: str, *, stata_rc: int | None = None):
        super().__init__(message)
        self.message = message
        self.stata_rc = stata_rc


def handle_page_request(manager: UIChannelManager, body: dict[str, Any], *, view_id: str | None) -> dict[str, Any]:
    max_limit, max_vars, max_chars, _ = manager.limits()

    if view_id is None:
        dataset_id = str(body.get("datasetId", ""))
        frame = str(body.get("frame", "default"))
    else:
        view = manager.get_view(view_id)
        if view is None:
            raise HTTPError(404, "not_found", "View not found")
        dataset_id = view.dataset_id
        frame = view.frame

    # Parse offset (default 0 is valid since offset >= 0)
    try:
        offset = int(body.get("offset") or 0)
    except (ValueError, TypeError) as e:
        raise HTTPError(400, "invalid_request", f"offset must be a valid integer, got: {body.get('offset')!r}")

    # Parse limit (no default - must be explicitly provided)
    limit_raw = body.get("limit")
    if limit_raw is None:
        raise HTTPError(400, "invalid_request", "limit is required")
    try:
        limit = int(limit_raw)
    except (ValueError, TypeError) as e:
        raise HTTPError(400, "invalid_request", f"limit must be a valid integer, got: {limit_raw!r}")

    vars_req = body.get("vars", [])
    include_obs_no = bool(body.get("includeObsNo", False))

    # Parse sortBy parameter
    sort_by = body.get("sortBy", [])
    if sort_by is not None and not isinstance(sort_by, list):
        raise HTTPError(400, "invalid_request", f"sortBy must be an array, got: {type(sort_by).__name__}")
    if sort_by and not all(isinstance(s, str) for s in sort_by):
        raise HTTPError(400, "invalid_request", "sortBy must be an array of strings")

    # Parse maxChars
    max_chars_raw = body.get("maxChars", max_chars)
    try:
        max_chars_req = int(max_chars_raw or max_chars)
    except (ValueError, TypeError) as e:
        raise HTTPError(400, "invalid_request", f"maxChars must be a valid integer, got: {max_chars_raw!r}")

    if offset < 0:
        raise HTTPError(400, "invalid_request", f"offset must be >= 0, got: {offset}")
    if limit <= 0:
        raise HTTPError(400, "invalid_request", f"limit must be > 0, got: {limit}")
    if limit > max_limit:
        raise HTTPError(400, "request_too_large", f"limit must be <= {max_limit}")
    if max_chars_req <= 0:
        raise HTTPError(400, "invalid_request", "maxChars must be > 0")
    if max_chars_req > max_chars:
        raise HTTPError(400, "request_too_large", f"maxChars must be <= {max_chars}")

    if not isinstance(vars_req, list) or not all(isinstance(v, str) for v in vars_req):
        raise HTTPError(400, "invalid_request", "vars must be a list of strings")
    if len(vars_req) > max_vars:
        raise HTTPError(400, "request_too_large", f"vars length must be <= {max_vars}")

    current_id = manager.current_dataset_id()
    if dataset_id != current_id:
        raise HTTPError(409, "dataset_changed", "Dataset changed")

    if view_id is None:
        obs_indices = None
        filtered_n: int | None = None
    else:
        assert view is not None
        obs_indices = view.obs_indices
        filtered_n = view.filtered_n

    try:
        # Apply sorting if requested
        if sort_by:
            try:
                manager._client.apply_sort(sort_by)
                # If sorting with a filtered view, re-compute indices after sort
                if view_id is not None:
                    assert view is not None
                    obs_indices = manager._client.compute_view_indices(view.filter_expr)
                    filtered_n = len(obs_indices)
            except ValueError as e:
                raise HTTPError(400, "invalid_request", f"Invalid sort specification: {e}")
            except RuntimeError as e:
                raise HTTPError(500, "internal_error", f"Failed to apply sort: {e}")

        dataset_state = manager._client.get_dataset_state()
        page = manager._client.get_page(
            offset=offset,
            limit=limit,
            vars=vars_req,
            include_obs_no=include_obs_no,
            max_chars=max_chars_req,
            obs_indices=obs_indices,
        )
    except HTTPError:
        # Re-raise HTTPError exceptions as-is
        raise
    except RuntimeError as e:
        # StataClient uses RuntimeError("No data in memory") for empty dataset.
        msg = str(e) or "No data in memory"
        if "no data" in msg.lower():
            raise HTTPError(400, "no_data_in_memory", msg)
        raise HTTPError(500, "internal_error", msg)
    except ValueError as e:
        msg = str(e)
        if msg.lower().startswith("invalid variable"):
            raise HTTPError(400, "invalid_variable", msg)
        raise HTTPError(400, "invalid_request", msg)
    except Exception as e:
        raise HTTPError(500, "internal_error", str(e))

    view_obj: dict[str, Any] = {
        "offset": offset,
        "limit": limit,
        "returned": page["returned"],
        "filteredN": filtered_n,
    }
    if view_id is not None:
        view_obj["viewId"] = view_id

    return {
        "dataset": {
            "id": current_id,
            "frame": dataset_state.get("frame"),
            "n": dataset_state.get("n"),
            "k": dataset_state.get("k"),
        },
        "view": view_obj,
        "vars": page["vars"],
        "rows": page["rows"],
        "display": {
            "maxChars": max_chars_req,
            "truncatedCells": page["truncated_cells"],
            "missing": ".",
        },
    }


def handle_arrow_request(manager: UIChannelManager, body: dict[str, Any], *, view_id: str | None) -> bytes:
    max_limit, max_vars, max_chars, _ = manager.limits()
    # Use the specific Arrow limit instead of the general UI page limit
    chunk_limit = getattr(manager, "_max_arrow_limit", 1_000_000)

    if view_id is None:
        dataset_id = str(body.get("datasetId", ""))
        frame = str(body.get("frame", "default"))
    else:
        view = manager.get_view(view_id)
        if view is None:
            raise HTTPError(404, "not_found", "View not found")
        dataset_id = view.dataset_id
        frame = view.frame

    # Parse offset (default 0)
    try:
        offset = int(body.get("offset") or 0)
    except (ValueError, TypeError):
        raise HTTPError(400, "invalid_request", "offset must be a valid integer")

    # Parse limit (required)
    limit_raw = body.get("limit")
    if limit_raw is None:
        # Default to the max arrow limit if not specified? 
        # The previous code required it. Let's keep it required but allow large values.
        raise HTTPError(400, "invalid_request", "limit is required")
    try:
        limit = int(limit_raw)
    except (ValueError, TypeError):
        raise HTTPError(400, "invalid_request", "limit must be a valid integer")

    vars_req = body.get("vars", [])
    include_obs_no = bool(body.get("includeObsNo", False))
    sort_by = body.get("sortBy", [])

    if offset < 0:
        raise HTTPError(400, "invalid_request", "offset must be >= 0")
    if limit <= 0:
        raise HTTPError(400, "invalid_request", "limit must be > 0")
    # Arrow streams are efficient, but we still respect a (much larger) max limit
    if limit > chunk_limit:
        raise HTTPError(400, "request_too_large", f"limit must be <= {chunk_limit}")

    if not isinstance(vars_req, list) or not all(isinstance(v, str) for v in vars_req):
        raise HTTPError(400, "invalid_request", "vars must be a list of strings")
    if len(vars_req) > max_vars:
        raise HTTPError(400, "request_too_large", f"vars length must be <= {max_vars}")

    current_id = manager.current_dataset_id()
    if dataset_id != current_id:
        raise HTTPError(409, "dataset_changed", "Dataset changed")

    if view_id is None:
        obs_indices = None
    else:
        assert view is not None
        obs_indices = view.obs_indices

    try:
        # Apply sorting if requested
        if sort_by:
            if not isinstance(sort_by, list) or not all(isinstance(s, str) for s in sort_by):
                raise HTTPError(400, "invalid_request", "sortBy must be a list of strings")
            try:
                manager._client.apply_sort(sort_by)
                if view_id is not None:
                    # encapsulated re-computation if view is active
                    # Note: original code only does this for view_id is not None
                    # But if we sort global dataset, existing views might become invalid unless
                    # they rely on stable indices. Stata indices change on sort.
                    # The current implementation of create_view computes indices once.
                    # If we sort, those indices point to different rows! 
                    # The original code handles this by re-computing view indices on sort.
                    assert view is not None
                    obs_indices = manager._client.compute_view_indices(view.filter_expr)
            except ValueError as e:
                raise HTTPError(400, "invalid_request", f"Invalid sort: {e}")
            except RuntimeError as e:
                raise HTTPError(500, "internal_error", f"Sort failed: {e}")

        arrow_bytes = manager._client.get_arrow_stream(
            offset=offset,
            limit=limit,
            vars=vars_req,
            include_obs_no=include_obs_no,
            obs_indices=obs_indices,
        )
        return arrow_bytes

    except RuntimeError as e:
        msg = str(e) or "No data in memory"
        if "no data" in msg.lower():
            raise HTTPError(400, "no_data_in_memory", msg)
        raise HTTPError(500, "internal_error", msg)
    except ValueError as e:
        msg = str(e)
        if "invalid variable" in msg.lower():
            raise HTTPError(400, "invalid_variable", msg)
        raise HTTPError(400, "invalid_request", msg)
    except Exception as e:
        raise HTTPError(500, "internal_error", str(e))

