services/unified_query.py

from __future__ import annotations

import re
from datetime import datetime
from typing import Any

OCEAN_VARIABLES = {
    "sst": ("sst", "海表温度", "sea surface temperature"),
    "ssh": ("ssh", "sla", "adt", "海表高度", "海平面高度", "sea surface height"),
    "chl": ("chl", "chlorophyll", "叶绿素"),
    "do": ("do", "oxygen", "溶解氧"),
    "u": ("uo", "u10", "u流", "纬向流", "东向流"),
    "v": ("vo", "v10", "v流", "经向流", "北向流"),
    "salinity": ("salinity", "so", "盐度"),
    "temperature_3d": ("thetao", "温度剖面", "三维温度"),
    "mld": ("mld", "mixed layer", "混合层"),
}

TUNA_TERMS = (
    "tuna","金枪鱼","鲣","黄鳍","大眼","长鳍","蓝鳍",
    "wcpfc","iotc","iccat","iattc"
)
SQUID_TERMS = (
    "squid","柔鱼","鱿鱼","茎柔鱼","美洲大赤鱿",
    "npfc","sprfmo"
)
FISHERY_TERMS = (
    "cpue","catch","effort","捕捞","渔业","渔获","渔具","作业","渔场"
)

def _uniq(items):
    out=[]
    seen=set()
    for x in items:
        if x not in seen:
            seen.add(x);out.append(x)
    return out

def _extract_date_range(q: str) -> dict[str, str]:
    q=q or ""
    iso=re.findall(r"\b(19\d{2}|20\d{2})[-/.](0?[1-9]|1[0-2])(?:[-/.](0?[1-9]|[12]\d|3[01]))?\b",q)
    zh=re.findall(r"(19\d{2}|20\d{2})\s*年\s*(0?[1-9]|1[0-2])\s*月(?:\s*(0?[1-9]|[12]\d|3[01])\s*日)?",q)
    years=re.findall(r"\b(19\d{2}|20\d{2})\b",q)
    result={}
    if zh:
        y,m,d=zh[0]
        result["start"]=f"{y}-{int(m):02d}-{int(d or 1):02d}"
        if len(zh)>1:
            y2,m2,d2=zh[-1]
            result["end"]=f"{y2}-{int(m2):02d}-{int(d2 or 28):02d}"
        else:
            result["end"]=result["start"] if d else f"{y}-{int(m):02d}-28"
    elif iso:
        y,m,d=iso[0]
        result["start"]=f"{y}-{int(m):02d}-{int(d or 1):02d}"
        if len(iso)>1:
            y2,m2,d2=iso[-1]
            result["end"]=f"{y2}-{int(m2):02d}-{int(d2 or 28):02d}"
        else:
            result["end"]=result["start"] if d else f"{y}-{int(m):02d}-28"
    elif years:
        result["start"]=f"{years[0]}-01-01"
        result["end"]=f"{years[-1]}-12-31"
    return result

def _extract_bbox(q: str) -> dict[str, float]:
    # Conservative parser for explicit E/W/N/S coordinate forms.
    q=(q or "").upper().replace("°","")
    lon=re.findall(r"(-?\d+(?:\.\d+)?)\s*([EW])",q)
    lat=re.findall(r"(-?\d+(?:\.\d+)?)\s*([NS])",q)
    def conv(pair):
        v=float(pair[0]); hemi=pair[1]
        if hemi in ("W","S"): v=-abs(v)
        else: v=abs(v)
        return v
    if len(lon)>=2 and len(lat)>=2:
        xs=[conv(x) for x in lon[:2]]
        ys=[conv(x) for x in lat[:2]]
        return {"west":min(xs),"east":max(xs),"south":min(ys),"north":max(ys)}
    return {}

def plan_unified_query(prompt: str) -> dict[str, Any]:
    q=str(prompt or "").strip()
    low=q.lower()

    variables=[]
    for canonical,terms in OCEAN_VARIABLES.items():
        if any(term.lower() in low for term in terms):
            variables.append(canonical)

    tuna=any(t.lower() in low for t in TUNA_TERMS)
    squid=any(t.lower() in low for t in SQUID_TERMS)
    fishery= tuna or squid or any(t.lower() in low for t in FISHERY_TERMS)

    # Ocean should be selected when environmental variables, explicit ocean
    # terms, or marine physical-process language is present.
    ocean=bool(variables) or any(t in low for t in (
        "ocean","海洋","环境因子","海温","海流","涡旋","水团","锋面","海表","叶绿素"
    ))

    domains=[]
    if ocean: domains.append("ocean")
    if tuna: domains.append("tuna")
    if squid: domains.append("squid")
    if fishery and not (tuna or squid): domains.append("fisheries")

    # If user explicitly asks for a joint habitat/CPUE analysis but species
    # family is not clear, keep Fisheries generic instead of guessing Tuna/Squid.
    intent="general_query"
    if ocean and fishery:
        intent="joint_ocean_fisheries"
    elif ocean:
        intent="ocean_query"
    elif fishery:
        intent="fisheries_query"

    date_range=_extract_date_range(q)
    bbox=_extract_bbox(q)

    missing=[]
    if not date_range:
        missing.append("时间范围")
    if ocean and not bbox and not any(x in low for x in ("全球","global")):
        missing.append("经纬度范围或更具体海域")
    if fishery and not (tuna or squid):
        missing.append("目标物种或渔业数据类型")

    return {
        "prompt":q,
        "intent":intent,
        "domains":domains or ["unknown"],
        "variables":_uniq(variables),
        "date_range":date_range,
        "bbox":bbox,
        "species_family":"tuna" if tuna else ("squid" if squid else ""),
        "missing_conditions":missing,
        "can_execute_preview":bool(domains),
    }