from __future__ import annotations

import json
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Iterable, List, Mapping


@dataclass(frozen=True)
class StageConfig:
    name: str
    agent: str
    needs: list[str]
    produces: list[str]
    task_template: str
    description: str | None = None


@dataclass(frozen=True)
class StateField:
    key: str
    required: bool = False
    prompt: str | None = None  # hint for UserProxy when missing
    description: str | None = None


@dataclass(frozen=True)
class FieldGlossaryEntry:
    key: str
    description: str
    aliases: list[str] = field(default_factory=list)


@dataclass(frozen=True)
class WorkflowConfig:
    task_type: str
    stages: list[StageConfig]
    team_agents: dict[str, list[StageConfig]] = field(default_factory=dict)
    agent_tools: dict[str, list[str]] = field(default_factory=dict)
    state_keys: dict[str, Any] = field(default_factory=dict)
    state_fields: list[StateField] = field(default_factory=list)
    field_glossary: list[FieldGlossaryEntry] = field(default_factory=list)

    def summary_lines(self) -> list[str]:
        out: list[str] = []
        for s in self.stages:
            needs = ", ".join(s.needs)
            produces = ", ".join(s.produces)
            description = f" {s.description}" if s.description else ""
            # task_template = f" (template: {s.task_template})" if s.task_template else ""
            out.append(f"{s.name}: agent={s.agent};  {description}; needs=[{needs}]; produces=[{produces}];")
        return out

    def required_state_lines(self) -> list[str]:
        out: list[str] = []
        for f in self.state_fields:
            if f.required:
                suffix = f" (ask: {f.prompt})" if f.prompt else ""
                out.append(f"{f.key}{suffix}")
        return out

    def field_glossary_lines(self) -> list[str]:
        out: list[str] = []
        if self.field_glossary:
            for entry in self.field_glossary:
                alias = f" (aliases: {', '.join(entry.aliases)})" if entry.aliases else ""
                out.append(f"{entry.key}: {entry.description}{alias}")
            return out
        for f in self.state_fields:
            if f.description:
                out.append(f"{f.key}: {f.description}")
        return out


class WorkflowNotFoundError(FileNotFoundError):
    pass


def _to_stage(obj: Mapping[str, Any]) -> StageConfig:
    return StageConfig(
        name=str(obj["name"]),
        agent=str(obj["agent"]),
        needs=list(obj.get("needs", [])),
        produces=list(obj.get("produces", [])),
        task_template=str(obj.get("task_template", "")),
        description=str(obj.get("description")) if obj.get("description") is not None else None,
    )


def _to_state_field(obj: Mapping[str, Any]) -> StateField:
    return StateField(
        key=str(obj["key"]),
        required=bool(obj.get("required", False)),
        prompt=str(obj.get("prompt")) if obj.get("prompt") is not None else None,
        description=str(obj.get("description")) if obj.get("description") is not None else None,
    )


def _to_glossary_entry(obj: Mapping[str, Any]) -> FieldGlossaryEntry:
    aliases = obj.get("aliases", [])
    return FieldGlossaryEntry(
        key=str(obj["key"]),
        description=str(obj.get("description", "")),
        aliases=[str(x) for x in aliases] if isinstance(aliases, list) else [],
    )


def load_workflow(task_type: str, *, tasks_dir: Path | None = None) -> WorkflowConfig:
    base = tasks_dir or Path(__file__).resolve().parents[1] / "Task"
    path = base / f"{task_type}.json"
    if not path.exists():
        raise WorkflowNotFoundError(f"Workflow config not found for task_type={task_type!r} at {path}")
    data = json.loads(path.read_text(encoding="utf-8"))
    stages = [_to_stage(x) for x in data.get("stages", [])]
    raw_team_agents = data.get("team_agents", {})
    team_agents: dict[str, list[StageConfig]] = {}
    if isinstance(raw_team_agents, dict):
        for team_name, team_stages in raw_team_agents.items():
            if isinstance(team_stages, list):
                team_agents[str(team_name)] = [
                    _to_stage(x) for x in team_stages if isinstance(x, dict)
                ]
    agent_tools = {k: list(v) for k, v in data.get("agent_tools", {}).items()}
    state_keys = data.get("state_keys", {})
    state_fields = [_to_state_field(x) for x in data.get("state_fields", []) if isinstance(x, dict)]
    field_glossary = [
        _to_glossary_entry(x) for x in data.get("field_glossary", []) if isinstance(x, dict)
    ]
    return WorkflowConfig(
        task_type=str(data.get("task_type", task_type)),
        stages=stages,
        team_agents=team_agents,
        agent_tools=agent_tools,
        state_keys=state_keys if isinstance(state_keys, dict) else {},
        state_fields=state_fields,
        field_glossary=field_glossary,
    )


def format_workflow_for_prompt(workflow: WorkflowConfig) -> str:
    lines = workflow.summary_lines()
    return "Workflow stages (in order):\n- " + "\n- ".join(lines)
