from __future__ import annotations

from typing import Any, Dict, Iterable, List, Optional, Tuple
from pydantic import BaseModel, ConfigDict, Field, field_validator


class IOItem(BaseModel):
    """
    Generic item that can appear as:
      - {"id": "...", "io": "..."}  (PLC registers in examples)
      - {"name": "...", "io": "..."} (your generated HIL physical_values)
    """
    model_config = ConfigDict(extra="allow")

    id: Optional[str] = None
    name: Optional[str] = None
    io: Optional[str] = None

    @property
    def key(self) -> Optional[str]:
        return self.id or self.name

    @field_validator("io")
    @classmethod
    def _validate_io(cls, v: Optional[str]) -> Optional[str]:
        if v is None:
            return v
        if v not in ("input", "output"):
            raise ValueError("io must be 'input' or 'output'")
        return v


def _iter_io_items(node: Any) -> Iterable[IOItem]:
    """
    Flatten nested structures into IOItem objects.
    Supports:
      - list[dict]
      - dict[str, list[dict]] (register groups)
      - dict[str, dict] where key is the id/name
    """
    if node is None:
        return

    if isinstance(node, list):
        for it in node:
            yield from _iter_io_items(it)
        return

    if isinstance(node, dict):
        # If it's directly a single IO item
        if ("id" in node or "name" in node) and "io" in node:
            yield IOItem.model_validate(node)
            return

        # Mapping form: {"<signal>": {"io": "...", ...}}
        for k, v in node.items():
            if isinstance(k, str) and isinstance(v, dict) and "io" in v and not ("id" in v or "name" in v):
                tmp = dict(v)
                tmp["id"] = k
                yield IOItem.model_validate(tmp)
            else:
                yield from _iter_io_items(v)
        return

    return


class PLC(BaseModel):
    model_config = ConfigDict(extra="allow")

    name: Optional[str] = None
    id: Optional[str] = None
    logic: str
    registers: Any = None  # can be dict of groups or list depending on source JSON

    @property
    def label(self) -> str:
        return str(self.id or self.name or "plc")

    def io_ids(self) -> Tuple[List[str], List[str]]:
        ins: List[str] = []
        outs: List[str] = []
        for item in _iter_io_items(self.registers):
            if item.io == "input" and item.key:
                ins.append(item.key)
            elif item.io == "output" and item.key:
                outs.append(item.key)
        return ins, outs


class HIL(BaseModel):
    model_config = ConfigDict(extra="allow")

    name: Optional[str] = None
    id: Optional[str] = None
    logic: str
    physical_values: Any = None  # list or dict depending on source JSON

    @property
    def label(self) -> str:
        return str(self.id or self.name or "hil")

    def pv_io(self) -> Tuple[List[str], List[str]]:
        ins: List[str] = []
        outs: List[str] = []
        for item in _iter_io_items(self.physical_values):
            if item.io == "input" and item.key:
                ins.append(item.key)
            elif item.io == "output" and item.key:
                outs.append(item.key)
        return ins, outs


class Config(BaseModel):
    """
    MVP config: we only care about PLCs and HILs for logic generation.
    Keep extra='allow' so future keys don't break parsing.
    """
    model_config = ConfigDict(extra="allow")

    plcs: List[PLC] = Field(default_factory=list)
    hils: List[HIL] = Field(default_factory=list)