from __future__ import annotations

from typing import Dict, List, Any, Union, TypeAlias, TypedDict, Annotated
from dataclasses import dataclass
from pydantic import BaseModel, Field
from .utils import merge_dicts, deep_merge

FieldSpec: TypeAlias = Union[str, Dict[str, List[Any]]]

@dataclass(frozen=True)
class GroupTask:
    path: tuple[str, ...]
    keys: tuple[str, ...]


class OverallState(TypedDict):
    fields: List[FieldSpec]
    tasks: List[GroupTask]

    memory: str
    memory_update_success: bool
    memory_update_resp: str
    memory_patch: Annotated[Dict[str, Any], deep_merge]

    final_result: Annotated[Dict[str, Any], deep_merge]
    context_by_task: Annotated[Dict[str, List[str]], merge_dicts]
    errors: Annotated[Dict[str, str], merge_dicts]
    

class TaskState(TypedDict, total=False):
    task: GroupTask
    query: str
    docs: List[str]
    memory_json: str
    memory_patch: Dict[str, Any]

    final_result: Dict[str, Any]
    context_by_task: Dict[str, List[str]]
    errors: Dict[str, str]



class WorkflowInitArgs(BaseModel):
    alayamem_url: str = Field(..., description="AlayaMem 服务地址")


class Input(BaseModel):
    fields: List[FieldSpec] = Field(..., description="要填写的表格字段模板（支持嵌套结构）")


class WorkflowContext(BaseModel):
    user_id: str = Field(..., description="用户ID")
    profile_id: str = Field(..., description="用户画像ID")
    extract_model_id: str = Field(..., description="用于抽取字段的模型ID")
    rerank_model_id: str = Field(default=None, description="用于重排序的模型ID（可选）")


class Output(BaseModel):
    final_result: Dict[str, Any] = Field(..., description="填写完成的表格数据（嵌套结构）")
    context_by_task: Dict[str, List[str]] = Field(default_factory=dict, description="每个任务检索到的文档片段（调试用）")
    errors: Dict[str, str] = Field(default_factory=dict, description="错误信息")
    tasks: List[GroupTask] = Field(default_factory=list, description="规划的任务列表")

