"""Read GCP Delta Lake tables by downloading parquet files from GCS and querying locally."""

from __future__ import annotations

import hashlib
import json
import os
import re
from pathlib import Path
from typing import Any
from urllib.parse import quote

import duckdb
import httpx

from src.utils.errors import DeltaError
from src.utils.logging_utils import get_logger

DEFAULT_LIMIT = 10_000
GCS_TIMEOUT = 30.0
MAX_PARQUET_FILES = 2_000
DEFAULT_CACHE_ROOT = "~/.cache/o11y-mcp-server/gcp-parquet"
_GS_URI_PATTERN = re.compile(r"^gs://([^/]+)/(.+)$")
_DATE_EQ_PATTERN = re.compile(
    r"\bdate\s*=\s*'(?P<date>\d{8})'",
    re.IGNORECASE,
)
_DATE_TERM_PATTERN = re.compile(r"^\s*date\s*=\s*'\d{8}'\s*$", re.IGNORECASE)
_CLAUSE_SPLIT_PATTERN = re.compile(r"\bORDER\s+BY\b|\bLIMIT\b", re.IGNORECASE)
_SAFE_QUERY_PATTERN = re.compile(
    r"^\s*SELECT\s+.+FROM\s+__t__(?:\s+WHERE\s+.+)?(?:\s+ORDER\s+BY\s+.+)?(?:\s+LIMIT\s+\d+)?\s*$",
    re.IGNORECASE | re.DOTALL,
)

logger = get_logger("gcp_parquet")


def _parse_gs_uri(table_path: str) -> tuple[str, str]:
    match = _GS_URI_PATTERN.match((table_path or "").strip())
    if not match:
        raise DeltaError("GCP table_path must be a gs:// URI with bucket and path.")
    bucket, prefix = match.group(1), match.group(2)
    if not bucket or not prefix:
        raise DeltaError("GCP table_path must include bucket and path.")
    return bucket, prefix.rstrip("/") + "/"


def _extract_single_date_partition(query: str | None) -> str:
    """
    Extract one date partition value from SQL filter: date='yyyymmdd'.
    """
    sql = (query or "").strip()
    if not sql:
        raise DeltaError(
            "GCP reads require a date partition filter in query, for example "
            "WHERE date='20260318'."
        )
    matches = [m.group("date") for m in _DATE_EQ_PATTERN.finditer(sql)]
    unique = list(dict.fromkeys(matches))
    if not unique:
        raise DeltaError(
            "GCP reads require date='yyyymmdd' in query to limit download scope."
        )
    if len(unique) > 1:
        raise DeltaError(
            "GCP reads support exactly one date partition per request. "
            f"Found multiple dates: {unique}"
        )
    return unique[0]


def _strip_partition_date_filter(query: str) -> str:
    """
    Remove date='yyyymmdd' term from WHERE clause for local parquet query execution.

    The date predicate is used to select the partition during download. Downloaded parquet
    files may not contain a materialized `date` column.
    """
    sql = query.strip()
    where_match = re.search(r"\bWHERE\b", sql, re.IGNORECASE)
    if not where_match:
        return sql

    head = sql[:where_match.start()]
    tail = sql[where_match.end():]
    clause_match = _CLAUSE_SPLIT_PATTERN.search(tail)
    if clause_match:
        where_body = tail[:clause_match.start()]
        suffix = tail[clause_match.start():]
    else:
        where_body = tail
        suffix = ""

    terms = [t.strip() for t in re.split(r"(?i)\s+AND\s+", where_body) if t.strip()]
    if not terms:
        return sql
    filtered = [t for t in terms if not _DATE_TERM_PATTERN.match(t)]
    if len(filtered) == len(terms):
        return sql

    if not filtered:
        rebuilt = f"{head.strip()} {suffix.strip()}".strip()
        return re.sub(r"\s+", " ", rebuilt)

    rebuilt = f"{head}WHERE {' AND '.join(filtered)} {suffix}".strip()
    return re.sub(r"\s+", " ", rebuilt)


def _list_parquet_objects(
    client: httpx.Client,
    bucket: str,
    prefix: str,
    access_token: str,
) -> list[str]:
    headers = {"Authorization": f"Bearer {access_token}"}
    base_url = f"https://storage.googleapis.com/storage/v1/b/{bucket}/o"
    names: list[str] = []
    page_token: str | None = None

    while True:
        params = {
            "prefix": prefix,
            "fields": "items(name),nextPageToken",
        }
        if page_token:
            params["pageToken"] = page_token
        resp = client.get(base_url, headers=headers, params=params)
        if resp.status_code >= 400:
            raise DeltaError(f"Failed to list GCS objects: HTTP {resp.status_code}")
        payload = resp.json()
        for item in payload.get("items", []) or []:
            name = item.get("name")
            if isinstance(name, str) and name.endswith(".parquet"):
                names.append(name)
                if len(names) >= MAX_PARQUET_FILES:
                    return names
        page_token = payload.get("nextPageToken")
        if not page_token:
            break

    return names


def _download_gcs_object(
    client: httpx.Client,
    bucket: str,
    name: str,
    access_token: str,
    dest_file: Path,
) -> None:
    headers = {"Authorization": f"Bearer {access_token}"}
    encoded_name = quote(name, safe="")
    url = f"https://storage.googleapis.com/storage/v1/b/{bucket}/o/{encoded_name}"
    resp = client.get(url, headers=headers, params={"alt": "media"})
    if resp.status_code >= 400:
        raise DeltaError(f"Failed to download GCS object {name}: HTTP {resp.status_code}")
    dest_file.write_bytes(resp.content)


def _cache_root() -> Path:
    raw = os.getenv("O11Y_GCP_PARQUET_CACHE_DIR", DEFAULT_CACHE_ROOT)
    return Path(raw).expanduser().resolve()


def _cache_dir_for_partition(bucket: str, prefix: str, day: str) -> Path:
    key = f"{bucket}|{prefix}|date={day}"
    digest = hashlib.sha256(key.encode("utf-8")).hexdigest()[:16]
    return _cache_root() / digest


def _cache_manifest(cache_dir: Path) -> Path:
    return cache_dir / "_cache_manifest.json"


def _is_cache_ready(cache_dir: Path) -> bool:
    manifest = _cache_manifest(cache_dir)
    if not cache_dir.exists() or not manifest.exists():
        return False
    try:
        meta = json.loads(manifest.read_text(encoding="utf-8"))
        expected = int(meta.get("parquet_count", 0))
    except Exception:
        return False
    actual = len(list(cache_dir.glob("*.parquet")))
    return expected > 0 and actual >= expected


def _write_cache_manifest(
    cache_dir: Path,
    *,
    bucket: str,
    prefix: str,
    day: str,
    parquet_count: int,
) -> None:
    data = {
        "bucket": bucket,
        "prefix": prefix,
        "date": day,
        "parquet_count": parquet_count,
    }
    _cache_manifest(cache_dir).write_text(
        json.dumps(data, indent=2, ensure_ascii=False),
        encoding="utf-8",
    )


def _ensure_partition_cache(
    *,
    client: httpx.Client,
    bucket: str,
    prefix: str,
    day: str,
    access_token: str,
) -> Path:
    cache_dir = _cache_dir_for_partition(bucket, prefix, day)
    if _is_cache_ready(cache_dir):
        logger.info("gcp_parquet_cache_hit cache_dir=%s", cache_dir)
        return cache_dir

    cache_dir.mkdir(parents=True, exist_ok=True)
    for file in cache_dir.glob("*.parquet"):
        file.unlink()

    partition_prefix = f"{prefix}date={day}/"
    parquet_names = _list_parquet_objects(client, bucket, partition_prefix, access_token)
    if not parquet_names:
        logger.info(
            "gcp_parquet_no_objects bucket=%s partition_prefix=%s cache_dir=%s",
            bucket,
            partition_prefix,
            cache_dir,
        )
        return cache_dir

    logger.info(
        "gcp_parquet_download_start bucket=%s partition_prefix=%s parquet_count=%s cache_dir=%s",
        bucket,
        partition_prefix,
        len(parquet_names),
        cache_dir,
    )
    for idx, name in enumerate(parquet_names):
        local_file = cache_dir / f"part_{idx}.parquet"
        _download_gcs_object(client, bucket, name, access_token, local_file)
        logger.info(
            "gcp_parquet_downloaded index=%s total=%s object=%s local_file=%s",
            idx + 1,
            len(parquet_names),
            name,
            local_file,
        )
    _write_cache_manifest(
        cache_dir,
        bucket=bucket,
        prefix=prefix,
        day=day,
        parquet_count=len(parquet_names),
    )
    return cache_dir


def read_gcp_delta_parquet(
    table_path: str,
    access_token: str,
    query: str | None = None,
    limit: int | None = None,
) -> dict[str, Any]:
    # Enforce single-day partition scope for GCP to avoid downloading the full table.
    token = (access_token or "").strip()
    if not token:
        raise DeltaError("GCP access_token is required for GCS parquet reads.")

    bucket, prefix = _parse_gs_uri(table_path)
    day = _extract_single_date_partition(query)
    with httpx.Client(timeout=GCS_TIMEOUT) as client:
        cache_dir = _ensure_partition_cache(
            client=client,
            bucket=bucket,
            prefix=prefix,
            day=day,
            access_token=token,
        )
        local_cache_path = f"file://{cache_dir}"
        if not _is_cache_ready(cache_dir):
            return {"rows": [], "local_cache_path": local_cache_path}

        glob_pattern = str(cache_dir / "*.parquet")
        conn = duckdb.connect(":memory:")
        try:
            # Some partitions may contain duplicated parquet rows. Build __t__ on a
            # deduplicated projection so downstream queries do not return duplicates.
            create_view_sql = (
                "CREATE VIEW __t__ AS SELECT DISTINCT * "
                f"FROM read_parquet('{glob_pattern}', union_by_name=true)"
            )
            conn.execute(create_view_sql)
            effective_limit = limit if limit is not None else DEFAULT_LIMIT
            if query and query.strip():
                normalized_query = _strip_partition_date_filter(query.strip())
                if not _SAFE_QUERY_PATTERN.match(normalized_query):
                    raise DeltaError(
                        "Optional query must be of the form: SELECT ... FROM __t__ "
                        "[WHERE ...] [ORDER BY ...] [LIMIT n]"
                    )
                sql = normalized_query
            else:
                sql = f"SELECT * FROM __t__ LIMIT {effective_limit}"

            try:
                result = conn.execute(sql)
                df = result.fetchdf()
            except Exception as e:
                raise DeltaError(f"DuckDB parquet read failed: {e}") from e
            if df is None or df.empty:
                return {"rows": [], "local_cache_path": local_cache_path}
            return {"rows": df.to_dict("records"), "local_cache_path": local_cache_path}
        finally:
            conn.close()