from __future__ import annotations

import asyncio
import calendar
import json
import re
import secrets
import time
from datetime import date, datetime, timedelta
from pathlib import Path
from typing import Any
from urllib.parse import urlparse

import httpx

VARIABLE_MAP={
    "sst":{"source":"oisst","variable":"sst","label":"SST"},
    "ssh":{"source":"cmems_surface","variable":"zos","label":"SSH"},
    "chl":{"source":"cmems_bgc","variable":"chl","label":"CHL"},
    "do":{"source":"cmems_bgc","variable":"o2","label":"DO"},
    "mld":{"source":"cmems_surface","variable":"mlotst","label":"MLD"},
    "thetao":{"source":"cmems_physics","variable":"thetao","label":"THETAO"},
    "uo":{"source":"cmems_physics","variable":"uo","label":"UO"},
    "vo":{"source":"cmems_physics","variable":"vo","label":"VO"},
}

# Named-region defaults used only when the user does not provide explicit
# coordinates. Keep these conservative and visible in the task metadata.
NAMED_OCEAN_REGIONS=(
    {
        "name":"西北太平洋",
        "aliases":("西北太平洋","northwest pacific","north-west pacific","nw pacific"),
        "bbox":{"lon_min":120.0,"lon_max":180.0,"lat_min":0.0,"lat_max":60.0},
    },
    {
        "name":"西北太平洋中纬度",
        "aliases":("西北太平洋中纬度","northwest pacific mid-latitude"),
        "bbox":{"lon_min":120.0,"lon_max":160.0,"lat_min":20.0,"lat_max":50.0},
    },
)

def _daterange(start:date,end:date):
    cur=start
    while cur<=end:
        yield cur
        cur+=timedelta(days=1)

def _parse_date_range(prompt:str):
    q=str(prompt or "")
    # Chinese month, e.g. 2002年8月
    m=re.search(r"(19\d{2}|20\d{2})\s*年\s*(0?[1-9]|1[0-2])\s*月",q)
    if m and not re.search(r"\d+\s*日",q):
        y,mo=int(m.group(1)),int(m.group(2))
        last=calendar.monthrange(y,mo)[1]
        return date(y,mo,1),date(y,mo,last)
    # Explicit ISO date range
    ds=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)
    if len(ds)>=2:
        a=date(*map(int,ds[0]));b=date(*map(int,ds[-1]))
        return (a,b) if a<=b else (b,a)
    return None

def _parse_bbox(prompt:str):
    raw=str(prompt or "")
    q=raw.upper().replace("°","")
    # Explicit coordinates always win over named-region defaults.
    q=q.replace("–","-").replace("—","-").replace("至","-")
    mlon=re.search(r"(-?\d+(?:\.\d+)?)\s*E\s*-\s*(-?\d+(?:\.\d+)?)\s*E",q)
    mlat=re.search(r"(-?\d+(?:\.\d+)?)\s*N\s*-\s*(-?\d+(?:\.\d+)?)\s*N",q)
    if mlon and mlat:
        xs=sorted((float(mlon.group(1)),float(mlon.group(2))))
        ys=sorted((float(mlat.group(1)),float(mlat.group(2))))
        return {
            "bbox":{"lon_min":xs[0],"lon_max":xs[1],"lat_min":ys[0],"lat_max":ys[1]},
            "region_name":"自定义经纬度范围",
            "region_source":"explicit_coordinates",
        }

    low=raw.lower()
    for item in NAMED_OCEAN_REGIONS:
        if any(alias.lower() in low for alias in item["aliases"]):
            return {
                "bbox":dict(item["bbox"]),
                "region_name":item["name"],
                "region_source":"named_region_default",
            }
    return None

def _parse_variables(prompt:str):
    low=str(prompt or "").lower()
    found=[]
    aliases={
        "sst":("sst","海表温度","海温"),
        "ssh":("ssh","海面高度","海表高度","zos"),
        "chl":("chl","叶绿素","chlorophyll"),
        "do":("溶解氧"," dissolved oxygen "," do "),
        "mld":("mld","混合层"),
        "thetao":("thetao","三维温度","温度剖面"),
        "uo":(" uo ","纬向流","东向流"),
        "vo":(" vo ","经向流","北向流"),
    }
    padded=" "+low+" "
    for key,terms in aliases.items():
        if any(t.lower() in padded for t in terms):
            found.append(key)
    return found

def parse_ocean_batch_request(prompt:str) -> dict[str,Any] | None:
    q=str(prompt or "").strip()
    low=q.lower()
    if not any(x in low for x in ("导出","下载","打包","export","download")):
        return None
    dr=_parse_date_range(q)
    region=_parse_bbox(q)
    variables=_parse_variables(q)
    if not dr or not region or not variables:
        return None
    bbox=region["bbox"]
    start,end=dr
    if start==end:
        return None
    days=(end-start).days+1
    if days>366:
        return None
    fmt="netcdf"
    if re.search(r"\bcsv\b",low):fmt="csv"
    elif re.search(r"\b(?:xlsx|excel)\b",low):fmt="xlsx"
    elif re.search(r"\bjson\b",low):fmt="json"
    return {
        "start_date":start.isoformat(),"end_date":end.isoformat(),
        "days":days,"variables":variables,"bbox":bbox,"format":fmt,
        "region_name":region.get("region_name") or "",
        "region_source":region.get("region_source") or "",
        "subtask_count":days*len(variables),
    }

class OceanBatchManager:
    def __init__(self,root:Path,api_url:str,concurrency:int=6):
        self.root=Path(root);self.root.mkdir(parents=True,exist_ok=True)
        self.api_url=str(api_url).rstrip("/")
        self.concurrency=max(1,min(int(concurrency),12))
        self.running:dict[str,asyncio.Task]={}
        # Jobs interrupted by process restart are resumable, not falsely running.
        for p in self.root.glob("*.json"):
            try:
                job=json.loads(p.read_text(encoding="utf-8"))
                if job.get("status") in {"running","queued"}:
                    job["status"]="paused"
                    job["message"]="Space/re进程重启后任务已暂停，可点击继续。"
                    self._write(job)
            except Exception:
                pass

    def _path(self,jid):return self.root/f"{jid}.json"
    def _read(self,jid):
        p=self._path(jid)
        if not p.exists():return None
        try:return json.loads(p.read_text(encoding="utf-8"))
        except Exception:return None
    def _write(self,job):
        job["updated_at"]=datetime.now().astimezone().isoformat(timespec="seconds")
        p=self._path(job["job_id"]);tmp=p.with_suffix(".json.tmp")
        tmp.write_text(json.dumps(job,ensure_ascii=False,indent=2),encoding="utf-8");tmp.replace(p)
        return job

    def _public(self,job):
        subs=job.get("subtasks") or []
        counts={k:0 for k in ("pending","running","completed","failed","canceled")}
        for s in subs:
            counts[s.get("status") if s.get("status") in counts else "pending"]+=1
        total=len(subs);done=counts["completed"]+counts["failed"]+counts["canceled"]
        return {
            "job_id":job["job_id"],"status":job.get("status"),"created_at":job.get("created_at"),
            "updated_at":job.get("updated_at"),"thread_id":job.get("thread_id", ""),"prompt":job.get("prompt"),"spec":job.get("spec"),
            "counts":counts,"total":total,"finished":done,
            "progress_percent":(done/total*100 if total else 0),
            "message":job.get("message",""),"cancel_requested":bool(job.get("cancel_requested")),
            "results":[
                {k:s.get(k) for k in ("subtask_id","date","label","source","variable","status","attempts","download_url","download_url_issued_at","filename","size_bytes","detail")}
                for s in subs if s.get("status") in {"completed","failed"}
            ],
            "known_size_bytes":sum(int(s.get("size_bytes") or 0) for s in subs if s.get("status")=="completed"),
            "known_size_file_count":sum(1 for s in subs if s.get("status")=="completed" and int(s.get("size_bytes") or 0)>0),
        }

    async def create(self,user_id:str,prompt:str,spec:dict[str,Any],thread_id:str=""):
        jid="ob_"+secrets.token_hex(10)
        start=date.fromisoformat(spec["start_date"]);end=date.fromisoformat(spec["end_date"])
        subs=[];n=0
        for d in _daterange(start,end):
            for key in spec["variables"]:
                meta=VARIABLE_MAP[key];n+=1
                subs.append({
                    "subtask_id":f"{n:04d}","date":d.isoformat(),"key":key,
                    "label":meta["label"],"source":meta["source"],"variable":meta["variable"],
                    "status":"pending","attempts":0,"download_url":"","filename":"","size_bytes":0,"detail":"","download_url_issued_at":"",
                })
        job={
            "job_id":jid,"user_id":str(user_id or ""),"thread_id":str(thread_id or ""),"prompt":prompt,"spec":spec,
            "status":"queued","cancel_requested":False,"message":"任务已创建，等待执行。",
            "created_at":datetime.now().astimezone().isoformat(timespec="seconds"),
            "subtasks":subs,
        }
        self._write(job)
        await self.start(jid)
        return self._public(self._read(jid))

    async def start(self,jid:str):
        existing=self.running.get(jid)
        if existing and not existing.done():return
        job=self._read(jid)
        if not job:return
        job["cancel_requested"]=False;job["status"]="running";job["message"]="正在并行执行单日 Ocean 导出子任务。"
        # failed subtasks are retried on resume; completed are preserved.
        for s in job.get("subtasks",[]):
            if s.get("status") in {"running","failed","canceled"}:s["status"]="pending"
        self._write(job)
        self.running[jid]=asyncio.create_task(self._run(jid))

    async def _one(self,jid:str,subtask_id:str,sem:asyncio.Semaphore):
        async with sem:
            job=self._read(jid)
            if not job or job.get("cancel_requested"):return
            sub=next((x for x in job["subtasks"] if x["subtask_id"]==subtask_id),None)
            if not sub or sub.get("status")=="completed":return
            sub["status"]="running";sub["attempts"]=int(sub.get("attempts") or 0)+1
            self._write(job)
            payload={
                "domain":"ocean","source":sub["source"],"date":sub["date"],"variable":sub["variable"],
                **job["spec"]["bbox"],"format":job["spec"].get("format") or "netcdf",
            }
            detail="";result={}
            for attempt in range(2):
                try:
                    async with httpx.AsyncClient(
                        timeout=httpx.Timeout(connect=8,read=120,write=20,pool=20),follow_redirects=True
                    ) as client:
                        r=await client.post(self.api_url+"/data/export",json=payload)
                        try:result=r.json()
                        except Exception:result={"detail":r.text[:500]}
                        if r.status_code<400 and isinstance(result,dict):
                            break
                        detail=str((result or {}).get("detail") or f"HTTP {r.status_code}")[:500]
                except Exception as exc:
                    detail=str(exc)[:500]
                if attempt==0:await asyncio.sleep(1.2)
            job=self._read(jid)
            if not job:return
            sub=next((x for x in job["subtasks"] if x["subtask_id"]==subtask_id),None)
            if not sub:return
            if isinstance(result,dict) and result.get("download_path") and not result.get("download_url"):
                result["download_url"]=self.api_url+str(result["download_path"])
            if isinstance(result,dict) and result.get("download_url"):
                sub["status"]="completed"
                sub["download_url"]=str(result["download_url"])
                sub["download_url_issued_at"]=datetime.now().astimezone().isoformat(timespec="seconds")
                guessed=Path(urlparse(sub["download_url"]).path).name
                ext={"netcdf":".nc","csv":".csv","xlsx":".xlsx","json":".json"}.get(job["spec"].get("format") or "netcdf",".bin")
                sub["filename"]=str(
                    result.get("filename")
                    or result.get("file_name")
                    or guessed
                    or f"{sub['source']}_{sub['variable']}_{sub['date'].replace('-','')}{ext}"
                )[:240]
                raw_size=(
                    result.get("size_bytes")
                    or result.get("file_size_bytes")
                    or result.get("content_length")
                    or result.get("bytes")
                    or 0
                )
                try:sub["size_bytes"]=max(0,int(raw_size))
                except Exception:sub["size_bytes"]=0
                sub["detail"]=""
            else:
                sub["status"]="failed";sub["detail"]=detail or str((result or {}).get("detail") or "导出失败")[:500]
            self._write(job)

    async def _run(self,jid:str):
        try:
            job=self._read(jid)
            if not job:return
            sem=asyncio.Semaphore(self.concurrency)
            ids=[s["subtask_id"] for s in job["subtasks"] if s.get("status")!="completed"]
            tasks=[asyncio.create_task(self._one(jid,sid,sem)) for sid in ids]
            if tasks:await asyncio.gather(*tasks,return_exceptions=True)
            job=self._read(jid)
            if not job:return
            if job.get("cancel_requested"):
                job["status"]="canceled";job["message"]="任务已取消。"
            else:
                failed=sum(1 for s in job["subtasks"] if s.get("status")=="failed")
                completed=sum(1 for s in job["subtasks"] if s.get("status")=="completed")
                if failed:
                    job["status"]="partial";job["message"]=f"已完成 {completed} 个，失败 {failed} 个；可仅重试失败项。"
                else:
                    job["status"]="completed";job["message"]="所有单日导出子任务已完成。"
            self._write(job)
        finally:
            self.running.pop(jid,None)

    async def cancel(self,jid:str):
        job=self._read(jid)
        if not job:return None
        job["cancel_requested"]=True
        for s in job["subtasks"]:
            if s.get("status")=="pending":s["status"]="canceled"
        job["status"]="canceling";job["message"]="正在停止尚未开始的子任务。"
        self._write(job)
        return self._public(job)

    async def retry_subtask(self,jid:str,subtask_id:str):
        job=self._read(jid)
        if not job:return None
        sub=next((x for x in job.get("subtasks",[]) if x.get("subtask_id")==subtask_id),None)
        if not sub:return None
        if sub.get("status")=="completed":return self._public(job)
        sub["status"]="pending";sub["detail"]=""
        job["status"]="running";job["cancel_requested"]=False
        job["message"]=f"正在重试子任务 {subtask_id}（{sub.get('date','')} {sub.get('label','')}）。"
        self._write(job)
        asyncio.create_task(self._retry_one_and_finalize(jid,subtask_id))
        return self._public(self._read(jid))

    async def _retry_one_and_finalize(self,jid:str,subtask_id:str):
        try:
            await self._one(jid,subtask_id,asyncio.Semaphore(1))
            job=self._read(jid)
            if not job:return
            failed=sum(1 for s in job.get("subtasks",[]) if s.get("status")=="failed")
            pending=sum(1 for s in job.get("subtasks",[]) if s.get("status") in {"pending","running"})
            completed=sum(1 for s in job.get("subtasks",[]) if s.get("status")=="completed")
            if pending:job["status"]="running"
            elif failed:
                job["status"]="partial";job["message"]=f"已完成 {completed} 个，仍有 {failed} 个失败。"
            else:
                job["status"]="completed";job["message"]="所有单日导出子任务已完成。"
            self._write(job)
        except Exception as exc:
            job=self._read(jid)
            if job:
                job["status"]="partial";job["message"]=f"子任务重试异常：{str(exc)[:200]}";self._write(job)

    def list_user(self,user_id:str,limit:int=100):
        items=[]
        for p in self.root.glob("ob_*.json"):
            try:
                job=json.loads(p.read_text(encoding="utf-8"))
                if str(job.get("user_id") or "")!=str(user_id or ""):continue
                items.append(self._public(job))
            except Exception:continue
        items.sort(key=lambda x:str(x.get("updated_at") or x.get("created_at") or ""),reverse=True)
        return items[:max(1,min(int(limit),500))]

    async def retry_failed(self,jid:str):
        job=self._read(jid)
        if not job:return None
        for s in job["subtasks"]:
            if s.get("status") in {"failed","canceled"}:s["status"]="pending";s["detail"]=""
        self._write(job);await self.start(jid)
        return self._public(self._read(jid))

    async def refresh_download_url(self,jid:str,subtask_id:str):
        """Re-export one completed subtask to obtain a fresh signed URL."""
        job=self._read(jid)
        if not job:return None
        sub=next((x for x in job.get("subtasks",[]) if x.get("subtask_id")==subtask_id),None)
        if not sub:return None
        payload={
            "domain":"ocean","source":sub["source"],"date":sub["date"],"variable":sub["variable"],
            **job["spec"]["bbox"],"format":job["spec"].get("format") or "netcdf",
        }
        result={};detail=""
        for attempt in range(3):
            try:
                async with httpx.AsyncClient(
                    timeout=httpx.Timeout(connect=8,read=120,write=20,pool=20),follow_redirects=True
                ) as client:
                    r=await client.post(self.api_url+"/data/export",json=payload)
                    try:result=r.json()
                    except Exception:result={"detail":r.text[:500]}
                    if r.status_code<400 and isinstance(result,dict) and (result.get("download_url") or result.get("download_path")):
                        break
                    detail=str((result or {}).get("detail") or f"HTTP {r.status_code}")[:500]
            except Exception as exc:
                detail=str(exc)[:500]
            if attempt<2:await asyncio.sleep(1.0*(2**attempt))
        if isinstance(result,dict) and result.get("download_path") and not result.get("download_url"):
            result["download_url"]=self.api_url+str(result["download_path"])
        if not isinstance(result,dict) or not result.get("download_url"):
            raise RuntimeError(detail or "无法刷新 Ocean 下载链接")
        job=self._read(jid)
        if not job:return None
        sub=next((x for x in job.get("subtasks",[]) if x.get("subtask_id")==subtask_id),None)
        if not sub:return None
        sub["download_url"]=str(result["download_url"])
        sub["download_url_issued_at"]=datetime.now().astimezone().isoformat(timespec="seconds")
        guessed=Path(urlparse(sub["download_url"]).path).name
        if result.get("filename") or result.get("file_name") or guessed:
            sub["filename"]=str(result.get("filename") or result.get("file_name") or guessed)[:240]
        raw_size=result.get("size_bytes") or result.get("file_size_bytes") or result.get("content_length") or result.get("bytes") or 0
        try:sub["size_bytes"]=max(0,int(raw_size))
        except Exception:pass
        sub["detail"]=""
        self._write(job)
        return {
            "job_id":jid,"subtask_id":subtask_id,"download_url":sub["download_url"],
            "download_url_issued_at":sub["download_url_issued_at"],
            "filename":sub.get("filename") or "download","size_bytes":int(sub.get("size_bytes") or 0),
        }

    def get(self,jid): 
        job=self._read(jid)
        return self._public(job) if job else None

    def manifest(self,jid):
        job=self._read(jid)
        if not job:return None
        files=[]
        for s in job["subtasks"]:
            if s.get("status")!="completed" or not s.get("download_url"):
                continue
            filename=s.get("filename") or Path(urlparse(s["download_url"]).path).name or (
                f"{s['source']}_{s['variable']}_{s['date'].replace('-','')}.nc"
            )
            files.append({
                "subtask_id":s["subtask_id"],
                "date":s["date"],
                "variable":s["label"],
                "source":s["source"],
                "filename":filename,
                "size_bytes":int(s.get("size_bytes") or 0),
                "download_url":s.get("download_url"),
                "download_url_issued_at":s.get("download_url_issued_at") or "",
            })
        return {
            "job_id":jid,"spec":job.get("spec"),"status":job.get("status"),
            "file_count":len(files),
            "known_size_bytes":sum(int(x.get("size_bytes") or 0) for x in files),
            "known_size_file_count":sum(1 for x in files if int(x.get("size_bytes") or 0)>0),
            "files":files,
            "failed":[
                {"date":s["date"],"variable":s["label"],"detail":s.get("detail")}
                for s in job["subtasks"] if s.get("status")=="failed"
            ],
        }