"""Token-efficient structured compression for data-request conversations.

The upstream LLM still owns the authoritative thread history.  This module
only builds a compact, structured context card for the parts of a data
request that must survive conversation trimming, thread expiry and follow-up
amendments:

    dataset/source, variables, region, depth, date, output format.

It never summarizes scientific values from the downloaded files and never
invents parameters that are not present in the user's own messages.
"""

from __future__ import annotations

import calendar
import copy
import re
from datetime import date
from typing import Any

from sidebar_catalog import (
    HF_SOURCE_ALIASES,
    OCEAN_CATALOG,
    OCEAN_VARIABLE_NAMES_ZH,
)

SOURCE_LABELS = {
    "era5": "ERA5",
    "era5_accum": "ERA5 Accumulation",
    "cmems_physics": "CMEMS Physics",
    "cmems_surface": "CMEMS Surface",
    "cmems_bgc": "CMEMS BGC",
    "cmems_carbonate": "CMEMS Carbonate",
    "occci": "OC-CCI",
    "oisst": "OISST",
}

SOURCE_ALIASES = (
    ("era5_accum", ("era5_accum", "era5累积", "era5 累积")),
    ("cmems_physics", ("cmems_physics", "cmems 物理", "物理场")),
    ("cmems_surface", ("cmems_surface", "cmems 海表")),
    ("cmems_bgc", ("cmems_bgc", "cmems 生物地球化学", "生物化学")),
    ("cmems_carbonate", ("cmems_carbonate", "cmems 碳酸盐")),
    ("era5", ("era5",)),
    ("cmems", ("cmems",)),
    ("occci", ("occci", "oc-cci", "海色")),
    ("oisst", ("oisst", "最优插值海表温度")),
)

VARIABLE_SOURCE_HINT = {}
for _src, _label, _label_zh, _vars in OCEAN_CATALOG:
    for _v in _vars:
        VARIABLE_SOURCE_HINT.setdefault(_v, _src)

VARIABLE_ALIASES = {}
for _code, _zh in OCEAN_VARIABLE_NAMES_ZH.items():
    VARIABLE_ALIASES[_zh] = _code

EXTRA_VARIABLE_ALIASES = {
    "海表温度": "sst",
    "海温": "sst",
    "海面高度": "zos",
    "海面高度异常": "zos",
    "叶绿素": "chl",
    "溶解氧": "o2",
    "混合层深度": "mlotst",
    "混合层": "mlotst",
    "盐度": "so",
    "海水温度": "thetao",
    "2米气温": "t2m",
    "2 米气温": "t2m",
    "10米风": "u10",
    "10 米风": "u10",
    "10米纬向风": "u10",
    "10米经向风": "v10",
    "总降水": "tp",
    "表层温度": "sst",
    "表层盐度": "so",
    "捕捞量": "catch",
    "捕捞": "catch",
    "努力量": "effort",
    "渔船活动": "effort",
    "资源量": "biomass",
    "资源评估": "biomass",
    "recruitment": "recruitment",
}
VARIABLE_ALIASES.update(EXTRA_VARIABLE_ALIASES)

FORMAT_ALIASES = (
    ("netcdf", ("netcdf", "netcdf4", r"\.nc\b", "nc格式")),
    ("csv", ("csv",)),
    ("xlsx", ("xlsx", "excel")),
    ("json", ("json",)),
    ("geotiff", ("geotiff", "tiff", "tif")),
    ("png", ("png",)),
    ("zip", ("zip", "打包", "压缩包")),
)

FISHERY_VARIABLES = {
    "catch": ("catch", "捕捞量", "捕捞"),
    "effort": ("effort", "努力量", "渔船活动"),
    "cpue": ("cpue",),
    "biomass": ("biomass", "资源量", "资源评估"),
    "recruitment": ("recruitment",),
    "fishing_hours": ("fishing hours", "表观捕捞时长", "apparent fishing"),
    "viirs": ("viirs", "夜光"),
}

DEPTH_RE = re.compile(
    r"(?:深度|depth)\s*[:：]?\s*"
    r"([0-9]+(?:\.[0-9]+)?(?:\s*(?:-|至|~)\s*[0-9]+(?:\.[0-9]+)?)?"
    r"\s*(?:米|m|dbar|db))",
    re.IGNORECASE,
)
DEPTH_BARE_RE = re.compile(
    r"(?<!\w)([0-9]+(?:\.[0-9]+)?(?:\s*(?:-|至|~)\s*[0-9]+(?:\.[0-9]+)?)?"
    r"\s*(?:米|m|dbar|db))(?!\w)",
    re.IGNORECASE,
)


def _clean_text(text: Any) -> str:
    return re.sub(r"\s+", " ", str(text or "")).strip()


def _norm(text: Any) -> str:
    return re.sub(
        r"[\s，,。.!！?？、；;:：()（）\[\]\"'“”‘’\-_/]",
        "",
        str(text or ""),
    ).lower()


def _dedupe(values: list[str]) -> list[str]:
    seen = set()
    out = []
    for value in values:
        key = str(value).strip().lower()
        if not key or key in seen:
            continue
        seen.add(key)
        out.append(str(value).strip())
    return out


def _looks_like_addition(text: str) -> bool:
    q = _norm(text)
    return any(
        token in q
        for token in ("加上", "增加", "新增", "添加", "补充", "同时", "另外", "还要")
    )


def extract_format(text: str) -> str:
    q = str(text or "").lower()
    for fmt, aliases in FORMAT_ALIASES:
        if any(re.search(alias, q) for alias in aliases):
            return fmt
    return ""


def _extract_sources(text: str) -> list[str]:
    q = str(text or "").lower()
    found = []
    for canonical, aliases in SOURCE_ALIASES:
        if canonical not in found and any(alias in q for alias in aliases):
            found.append(canonical)
    return found


def _extract_ocean_variables(text: str) -> list[str]:
    q = re.sub(r"[\s，,、。；;:/_]", " ", str(text or "").lower())
    q = " " + q + " "
    found = []
    # Chinese aliases are matched longest-first so “海表温度” wins over “温度”.
    for alias in sorted(VARIABLE_ALIASES, key=len, reverse=True):
        if alias in q:
            code = VARIABLE_ALIASES[alias]
            if code not in found:
                found.append(code)
    for code, _zh in OCEAN_VARIABLE_NAMES_ZH.items():
        if len(code) >= 2 and re.search(
            rf"(?<![a-z0-9_]){re.escape(code)}(?![a-z0-9_])",
            q,
        ):
            if code not in found:
                found.append(code)
    if "全部变量" in q or "所有变量" in q or "all variable" in q:
        found.append("ALL")
    return found


def _extract_fishery_variables(text: str) -> list[str]:
    q = _norm(text)
    found = []
    for code, aliases in FISHERY_VARIABLES.items():
        if any(alias in q for alias in aliases):
            found.append(code)
    return found


def _extract_fishery_sources(text: str) -> list[str]:
    q = _norm(text)
    found = []
    for source, aliases in HF_SOURCE_ALIASES.items():
        if any(alias in q for alias in aliases):
            found.append(source)
    for repo in ("squid_dataset", "tuna-fisheries-dataset"):
        if repo in q:
            found.append(repo)
    return found


def _extract_date_text(text: str) -> str:
    q = str(text or "")
    pieces = []
    date_part = (
        r"(?:19|20)\d{2}\s*年\s*(?:0?[1-9]|1[0-2])\s*月"
        r"(?:\s*(?:0?[1-9]|[12]\d|3[01])\s*日)?"
        r"|(?:19|20)\d{2}-(?:0[1-9]|1[0-2])-(?:0[1-9]|[12]\d|3[01])"
    )
    range_re = re.compile(
        "(" + date_part + r")"
        r"\s*(?:至|到|~|-|–|—)\s*"
        "(" + date_part + ")",
        re.IGNORECASE,
    )
    for match in range_re.finditer(q):
        pieces.append(
            f"{_normalize_date_part(match.group(1))} 至 "
            f"{_normalize_date_part(match.group(2))}"
        )
    q2 = range_re.sub("", q)

    # “2002年8月1日至8月31日” shorthand end date.
    shorthand_range = re.compile(
        r"((?:19|20)\d{2}\s*年\s*(?:0?[1-9]|1[0-2])\s*月"
        r"(?:\s*(?:0?[1-9]|[12]\d|3[01])\s*日)?)"
        r"\s*(?:至|到)\s*"
        r"((?:0?[1-9]|1[0-2])\s*月"
        r"(?:\s*(?:0?[1-9]|[12]\d|3[01])\s*日)?)",
        re.IGNORECASE,
    )
    for match in shorthand_range.finditer(q2):
        start = match.group(1)
        end = _fill_shorthand_date(start, match.group(2))
        pieces.append(f"{_normalize_date_part(start)} 至 {_normalize_date_part(end)}")
    q2 = shorthand_range.sub("", q2)

    single_re = re.compile(
        r"(?:19|20)\d{2}\s*年\s*(?:0?[1-9]|1[0-2])?\s*月?"
        r"(?:\s*(?:0?[1-9]|[12]\d|3[01])\s*日)?"
        r"|(?:19|20)\d{2}-(?:0?[1-9]|1[0-2])-(?:0?[1-9]|[12]\d|3[01])",
        re.IGNORECASE,
    )
    for match in single_re.finditer(q2):
        value = _normalize_date_part(match.group(0))
        if value not in pieces:
            pieces.append(value)
    return "、".join(_dedupe(pieces))


def _fill_shorthand_date(start: str, shorthand: str) -> str:
    year = re.search(r"((?:19|20)\d{2})\s*年", start)
    if not year or re.search(r"(?:19|20)\d{2}\s*年", shorthand):
        return shorthand
    return f"{year.group(1)}年{shorthand}"


def _normalize_date_part(value: str) -> str:
    value = str(value or "").strip()
    match = re.search(
        r"((?:19|20)\d{2})\s*年\s*(?:0?([1-9]|1[0-2]))?\s*月?"
        r"(?:\s*(?:0?([1-9]|[12]\d|3[01]))\s*日)?",
        value,
    )
    if match:
        year = match.group(1)
        month = match.group(2)
        day = match.group(3)
        return (
            f"{year}"
            + (f"-{int(month):02d}" if month else "")
            + (f"-{int(day):02d}" if day else "")
        )
    match = re.search(
        r"(19|20)\d{2}-(0?[1-9]|1[0-2])-(0?[1-9]|[12]\d|3[01])",
        value,
    )
    if match:
        return (
            f"{match.group(1)}-{int(match.group(2)):02d}-"
            f"{int(match.group(3)):02d}"
        )
    return value


def _extract_region_text(text: str) -> str:
    q = str(text or "")
    low = q.lower()
    # Explicit coordinate pairs with longitude and latitude.
    coord = re.search(
        r"(-?\d+(?:\.\d+)?)\s*[°]?\s*[eeww]\s*(?:-|至|–|—)\s*"
        r"(-?\d+(?:\.\d+)?)\s*[°]?\s*[eeww][^。；;\n]{0,30}?"
        r"(-?\d+(?:\.\d+)?)\s*[°]?\s*[nns]\s*(?:-|至|–|—)\s*"
        r"(-?\d+(?:\.\d+)?)\s*[°]?\s*[nns]",
        q,
        re.IGNORECASE,
    )
    if coord:
        lon1, lon2, lat1, lat2 = (float(x) for x in coord.groups())
        a, b = sorted((lon1, lon2))
        c, d = sorted((lat1, lat2))
        return f"{a:g}°E–{b:g}°E, {c:g}°N–{d:g}°N"

    # Partial longitude-only amendments are kept as raw text; an earlier
    # latitude is merged later when both halves exist.
    partial_lon = re.search(
        r"(-?\d+(?:\.\d+)?)\s*[°]?\s*[eeww]\s*(?:-|至|–|—)\s*"
        r"(-?\d+(?:\.\d+)?)\s*[°]?\s*[eeww]",
        q,
        re.IGNORECASE,
    )
    if partial_lon:
        return partial_lon.group(0).replace(" ", "")
    partial_lat = re.search(
        r"(-?\d+(?:\.\d+)?)\s*[°]?\s*[nns]\s*(?:-|至|–|—)\s*"
        r"(-?\d+(?:\.\d+)?)\s*[°]?\s*[nns]",
        q,
        re.IGNORECASE,
    )
    if partial_lat:
        return partial_lat.group(0).replace(" ", "")

    named_regions = (
        "西北太平洋",
        "西北太平洋中纬度",
        "太平洋",
        "大西洋",
        "印度洋",
        "南极",
        "全球",
        "南海",
        "东海",
        "黄海",
        "渤海",
    )
    for name in named_regions:
        if name in low or name.lower() in q:
            return name
    return ""


def extract_depth(text: str) -> str:
    q = str(text or "")
    match = DEPTH_RE.search(q)
    if match:
        return match.group(1).strip()
    match = DEPTH_BARE_RE.search(q)
    if match:
        return match.group(1).strip()
    if re.search(r"海表|表层|surface", q, re.IGNORECASE) and not re.search(
        r"cmems_surface|海表高度|surface height",
        q,
        re.IGNORECASE,
    ):
        return "surface"
    return ""


def extract_request_spec(text: Any) -> dict[str, Any]:
    """Extract the stable data-request fields from one user message."""
    source_text = _clean_text(text)
    q_low = str(text or "").lower()
    ocean_variables = _extract_ocean_variables(source_text)
    fishery_variables = _extract_fishery_variables(source_text)
    ocean_sources = _extract_sources(source_text)
    fishery_sources = _extract_fishery_sources(source_text)

    if fishery_sources or fishery_variables or any(
        term in q_low
        for term in (
            "柔鱼", "鱿鱼", "squid", "金枪鱼", "tuna", "fao", "gfw",
            "ram legacy", "sea around us", "捕捞", "努力量", "cpue",
            "biomass", "recruitment",
        )
    ):
        domain = "fisheries"
    elif ocean_sources or ocean_variables or any(
        term in q_low
        for term in ("ocean", "海洋", "海温", "sst", "era5", "cmems", "oisst")
    ):
        domain = "ocean"
    else:
        domain = ""

    variable_hint_source = ""
    for variable in ocean_variables:
        variable_hint_source = VARIABLE_SOURCE_HINT.get(variable, "")
        if variable_hint_source:
            break
    source_codes = ocean_sources or ([variable_hint_source] if variable_hint_source else [])
    if not source_codes and "oisst" in q_low:
        source_codes = ["oisst"]

    explicit_source = bool(ocean_sources or fishery_sources)
    dataset_origin = "explicit" if explicit_source else (
        "inferred" if source_codes else ""
    )
    dataset_parts = []
    for code in source_codes:
        label = SOURCE_LABELS.get(code, code)
        if label not in dataset_parts:
            dataset_parts.append(label)
    dataset_parts.extend(
        source for source in fishery_sources if source not in dataset_parts
    )

    variables = _dedupe(ocean_variables + fishery_variables)
    region = _extract_region_text(source_text)
    depth = extract_depth(source_text)
    date_text = _extract_date_text(source_text)
    output_format = extract_format(source_text)

    return {
        "domain": domain,
        "dataset": dataset_parts,
        "dataset_origin": dataset_origin,
        "variables": variables,
        "region": region,
        "depth": depth,
        "date": date_text,
        "format": output_format,
        "raw_request": source_text[:2000],
    }


def _merge_region(old_region: str, new_region: str) -> str:
    old_region = str(old_region or "")
    new_region = str(new_region or "")
    if not new_region:
        return old_region
    if not old_region:
        return new_region
    has_old_lon = bool(re.search(r"[eE]°|[eE]\b", old_region))
    has_old_lat = bool(re.search(r"[nN]°|[nN]\b", old_region))
    has_new_lon = bool(re.search(r"[eE]°|[eE]\b", new_region))
    has_new_lat = bool(re.search(r"[nN]°|[nN]\b", new_region))
    if has_old_lon and has_old_lat and has_new_lon and has_new_lat:
        return new_region
    if has_old_lat and has_new_lon and not has_new_lat:
        return new_region + ", " + old_region
    return new_region


def merge_request_spec(
    base: dict[str, Any] | None,
    text: Any,
    *,
    revision: bool = True,
) -> dict[str, Any]:
    """Merge a new user message into the running structured context card."""
    spec = copy.deepcopy(base) if isinstance(base, dict) else {}
    parsed = extract_request_spec(text)
    source_text = _clean_text(text)

    if parsed.get("domain"):
        spec["domain"] = parsed["domain"]
    existing_dataset_origin = spec.get("dataset_origin") or ""
    parsed_origin = parsed.get("dataset_origin") or ""
    if parsed.get("dataset"):
        if parsed_origin == "explicit":
            switch_marker = any(
                token in _norm(source_text)
                for token in ("换成", "换为", "改为", "改成", "改用", "替换")
            )
            if existing_dataset_origin == "explicit" and switch_marker:
                spec["dataset"] = parsed["dataset"]
            elif existing_dataset_origin == "explicit" and _looks_like_addition(source_text):
                spec["dataset"] = _dedupe(
                    list(spec.get("dataset") or []) + list(parsed["dataset"])
                )
            else:
                spec["dataset"] = parsed["dataset"]
            spec["dataset_origin"] = "explicit"
        elif existing_dataset_origin != "explicit" and parsed.get("dataset"):
            spec["dataset"] = parsed["dataset"]
            spec["dataset_origin"] = parsed_origin
    if parsed.get("variables"):
        existing_vars = list(spec.get("variables") or [])
        new_vars = list(parsed["variables"])
        if existing_vars and _looks_like_addition(source_text):
            spec["variables"] = _dedupe(existing_vars + new_vars)
        else:
            spec["variables"] = new_vars
    if parsed.get("region"):
        spec["region"] = _merge_region(spec.get("region", ""), parsed["region"])
    if parsed.get("depth"):
        spec["depth"] = parsed["depth"]
    if parsed.get("date"):
        spec["date"] = parsed["date"]
    if parsed.get("format"):
        spec["format"] = parsed["format"]

    # Keep only the first full user wording plus the latest amendment wording.
    # Older repeated messages are intentionally dropped to save tokens.
    if not spec.get("raw_request"):
        spec["raw_request"] = source_text[:1600]
    if revision:
        spec["last_revision"] = source_text[:900]
    elif source_text and spec.get("raw_request") == source_text[:1600]:
        spec["last_revision"] = ""
    spec["updated"] = True
    return spec


def render_request_spec(
    spec: dict[str, Any] | None,
    *,
    include_raw: bool = True,
) -> str:
    """Render the card as a small number of prompt tokens."""
    spec = spec if isinstance(spec, dict) else {}
    lines = []
    if spec.get("domain"):
        lines.append(f"数据域：{spec['domain']}")
    if spec.get("dataset"):
        lines.append(f"数据集/来源：{'、'.join(spec['dataset'])}")
    if spec.get("variables"):
        lines.append(f"变量：{'、'.join(spec['variables'])}")
    if spec.get("region"):
        lines.append(f"区域：{spec['region']}")
    if spec.get("depth"):
        lines.append(f"深度：{spec['depth']}")
    if spec.get("date"):
        lines.append(f"日期/时间：{spec['date']}")
    if spec.get("format"):
        lines.append(f"输出格式：{spec['format']}")
    if spec.get("last_revision"):
        lines.append(f"最近补充/修改：{spec['last_revision']}")
    has_structured = any(
        spec.get(key)
        for key in (
            "domain", "dataset", "variables", "region",
            "depth", "date", "format",
        )
    )
    if include_raw and spec.get("raw_request") and not has_structured:
        raw = str(spec["raw_request"]).replace("\n", " ")[:500]
        if raw not in str(spec.get("last_revision") or ""):
            lines.append(f"原始请求：{raw}")
    if not lines:
        return ""
    return (
        "[COMPRESSED_DATA_REQUEST_CONTEXT]\n"
        + "\n".join(lines)
        + "\n[/COMPRESSED_DATA_REQUEST_CONTEXT]"
    )


def compact_history(
    messages: list[dict[str, Any]] | None,
    *,
    max_recent_chars: int = 1400,
) -> str:
    """Build a compact context from stored messages without re-sending them."""
    messages = [m for m in (messages or []) if isinstance(m, dict)]
    if not messages:
        return ""

    spec = build_request_spec_from_messages(messages)

    rendered = render_request_spec(spec)

    recent = []
    total = 0
    recent_user_count = 0
    represented = {
        str(spec.get("raw_request") or "")[:1600],
        str(spec.get("last_revision") or ""),
    }
    for message in reversed(messages):
        role = str(message.get("role") or "")
        text = str(message.get("text") or "").strip()
        if role not in {"user", "assistant"} or not text:
            continue
        if role == "user" and (
            text[:1600] in represented
            or text == str(spec.get("raw_request") or "")
        ):
            continue
        prefix = "用户" if role == "user" else "助手"
        if role == "user":
            recent_user_count += 1
        snippet = f"{prefix}：{text}"
        if total + len(snippet) > max_recent_chars:
            snippet = snippet[: max(max_recent_chars - total, 1)]
            recent.append(snippet)
            break
        recent.append(snippet)
        total += len(snippet)
        if (
            recent_user_count >= 1
            and sum(1 for x in recent if x.startswith("助手")) >= 1
        ) or len(recent) >= 3:
            break
    recent_text = "\n".join(reversed(recent))

    parts = []
    if rendered:
        parts.append(rendered)
    if recent_text:
        parts.append("[RECENT_MESSAGES]\n" + recent_text + "\n[/RECENT_MESSAGES]")
    return "\n\n".join(parts)[:8000]


def build_request_spec_from_messages(
    messages: list[dict[str, Any]] | None,
) -> dict[str, Any]:
    """Merge all user messages that carry data parameters into one spec."""
    messages = [m for m in (messages or []) if isinstance(m, dict)]
    spec: dict[str, Any] = {}
    raw_found = False
    for message in messages:
        if str(message.get("role") or "") != "user":
            continue
        text = str(message.get("text") or "").strip()
        if not text:
            continue
        parsed = extract_request_spec(text)
        has_data = any(
            parsed.get(key)
            for key in (
                "domain", "dataset", "variables", "region",
                "depth", "date", "format",
            )
        )
        if has_data:
            if not raw_found:
                spec = merge_request_spec(spec, text, revision=False)
                spec["raw_request"] = text[:1600]
                raw_found = True
            else:
                spec = merge_request_spec(spec, text, revision=True)
    return spec