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 WorkflowInitArgs(BaseModel):
    alayamem_url: str = Field(..., description="AlayaMem 服务地址")


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


class WorkflowContext(BaseModel):
    collection_name: str = Field(default="file_watcher_collection", description="检索的集合名称")


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="规划的任务列表")


class OverallState(TypedDict):
    fields: List[FieldSpec]
    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]
    tasks: List[GroupTask]


class TaskState(TypedDict):
    task: GroupTask
