from __future__ import annotations
import json, time
from datetime import datetime
from pathlib import Path
from typing import Any, Callable

class UserStateStore:
    """Server-side persistence for workspace/settings/favorites.

    Keeps file persistence concerns out of the FastAPI assembly module.
    """

    def __init__(self, root: Path, validator: Callable[[str], bool], max_favorites: int = 300):
        self.root = Path(root)
        self.validator = validator
        self.max_favorites = int(max_favorites)

    def _require_user(self, user_id: str) -> str:
        uid = str(user_id or "").strip()
        if not self.validator(uid):
            raise ValueError("invalid user_id")
        return uid

    def workspace_state_file(self, user_id: str) -> Path:
        return self.root / f"{self._require_user(user_id)}.workspace.json"

    def sanitize_workspace_sessions(self, items: Any) -> list[dict[str, Any]]:
        out=[]
        for x in items if isinstance(items,list) else []:
            if not isinstance(x,dict): continue
            thread_id=str(x.get("thread_id") or "")[:200]
            if not thread_id: continue
            messages=[]
            for m in (x.get("messages") if isinstance(x.get("messages"),list) else [])[-80:]:
                if not isinstance(m,dict): continue
                messages.append({
                    "role":"assistant" if m.get("role")=="assistant" else "user",
                    "text":str(m.get("text") or "")[:50000],
                    "meta":str(m.get("meta") or "")[:1000],
                    "time":int(m.get("time") or int(time.time()*1000)),
                })
            out.append({
                "thread_id":thread_id,
                "title":str(x.get("title") or "")[:500],
                "messages":messages,
                "created_at":int(x.get("created_at") or int(time.time()*1000)),
                "updated_at":int(x.get("updated_at") or int(time.time()*1000)),
            })
        out.sort(key=lambda x:int(x.get("updated_at") or 0), reverse=True)
        return out[:100]

    def sanitize_workspace_settings(self, value: Any) -> dict[str, Any]:
        if not isinstance(value,dict): return {}
        allowed={"fontSize","enterToSend","refreshSeconds","accent","defaultView","compact","reduceMotion","longitudeRange","timeDisplay","exportFormat","answerDetail","historyLimit","confirmDownload","messageActionTrigger","sidebarToolsCollapsed"}
        return {str(k):v for k,v in value.items() if k in allowed}

    def read_workspace(self, user_id: str) -> dict[str, Any]:
        path=self.workspace_state_file(user_id)
        if not path.exists(): return {"sessions":[],"settings":{},"updated_at":""}
        try: raw=json.loads(path.read_text(encoding="utf-8"))
        except Exception: return {"sessions":[],"settings":{},"updated_at":""}
        return {
            "sessions":self.sanitize_workspace_sessions(raw.get("sessions",[])),
            "settings":self.sanitize_workspace_settings(raw.get("settings",{})),
            "updated_at":str(raw.get("updated_at") or ""),
        }

    def write_workspace(self, user_id: str, sessions: Any, settings: Any) -> dict[str, Any]:
        clean_sessions=self.sanitize_workspace_sessions(sessions)
        clean_settings=self.sanitize_workspace_settings(settings)
        payload={
            "schema":"squid-server-workspace",
            "schema_version":1,
            "updated_at":datetime.now().astimezone().isoformat(timespec="seconds"),
            "sessions":clean_sessions,
            "settings":clean_settings,
        }
        path=self.workspace_state_file(user_id)
        path.parent.mkdir(parents=True, exist_ok=True)
        tmp=path.with_suffix(path.suffix+".tmp")
        tmp.write_text(json.dumps(payload,ensure_ascii=False,indent=2),encoding="utf-8")
        tmp.replace(path)
        return payload

    def favorite_state_file(self, user_id: str) -> Path:
        return self.root / f"{self._require_user(user_id)}.favorites.json"

    def sanitize_favorite(self, item: Any) -> dict[str, Any] | None:
        if not isinstance(item, dict): return None
        kind=str(item.get("kind") or "answer")
        if kind not in {"dataset","file","answer"}: kind="answer"
        def cut(key: str, n: int) -> str: return str(item.get(key) or "")[:n]
        return {
            "kind":kind,
            "id":cut("id",240),
            "key":cut("key",800),
            "title":cut("title",400),
            "content":cut("content",60000),
            "prompt":cut("prompt",10000),
            "repository":cut("repository",500),
            "path":cut("path",2000),
            "created_at":int(item.get("created_at") or int(time.time()*1000)),
        }

    def read_favorites(self, user_id: str) -> list[dict[str, Any]]:
        path=self.favorite_state_file(user_id)
        if not path.exists(): return []
        try: raw=json.loads(path.read_text(encoding="utf-8"))
        except Exception: return []
        items=raw.get("favorites",[]) if isinstance(raw,dict) else raw
        out=[]
        for item in items if isinstance(items,list) else []:
            clean=self.sanitize_favorite(item)
            if clean: out.append(clean)
        return out[:self.max_favorites]

    def write_favorites(self, user_id: str, favorites: list[dict[str, Any]]) -> list[dict[str, Any]]:
        cleaned=[]; seen=set()
        for item in favorites:
            clean=self.sanitize_favorite(item)
            if not clean: continue
            ident=clean.get("id") or f"{clean.get('kind')}:{clean.get('key')}:{clean.get('title')}"
            if ident in seen: continue
            seen.add(ident); cleaned.append(clean)
            if len(cleaned)>=self.max_favorites: break
        path=self.favorite_state_file(user_id)
        path.parent.mkdir(parents=True, exist_ok=True)
        tmp=path.with_suffix(path.suffix+".tmp")
        payload={
            "schema":"squid-server-favorites",
            "schema_version":1,
            "updated_at":datetime.now().astimezone().isoformat(timespec="seconds"),
            "favorites":cleaned,
        }
        tmp.write_text(json.dumps(payload,ensure_ascii=False,indent=2),encoding="utf-8")
        tmp.replace(path)
        return cleaned

    def storage_mode(self) -> str:
        try:
            return "persistent" if str(self.root.resolve()).startswith("/data/") else "server_session"
        except Exception:
            return "server_session"