from __future__ import annotations
from typing import Any, Dict, List, Optional, Tuple, Type, TypeVar
from pydantic import BaseModel, Field as PydanticField
from logging import getLogger

logger = getLogger(__name__)

def _merge_json_extra(dst: Dict[str, Any], src: Dict[str, Any]) \
    -> Dict[str, Any]:

    out = dict(dst or {})
    for k, v in (src or {}).items():
        if k in out and isinstance(out[k], dict) and isinstance(v, dict):
            out[k] = {**out[k], **v}
        else:
            out[k] = v
    return out


def DirectField(
    *,
    target_name: str,
    targets: Tuple[Type[BaseModel], ...] | Type[BaseModel],
    src_name: Optional[str] = None,
    default: Any = ...,
    json_schema_extra: Optional[Dict[str, Any]] = None,
    **field_kwargs: Any,
):
    extra = _merge_json_extra(
        json_schema_extra or {},
        {
            "direction": {
                "target_name": target_name,
                "src_name": src_name,
                "targets": \
                    targets if isinstance(targets, tuple) else (targets,),
            }
        },
    )
    return Field(default, json_schema_extra=extra, **field_kwargs)


def HideField(
    default: Any = ...,
    *,
    json_schema_extra: Optional[Dict[str, Any]] = None,
    **field_kwargs: Any,
):
    extra = _merge_json_extra(json_schema_extra or {}, {"hide": True})
    return Field(default, json_schema_extra=extra, **field_kwargs)


def Field(
    default: Any = ...,
    *,
    target_name: Optional[str] = None,
    targets: Tuple[Type[BaseModel], ...] | Type[BaseModel] = (),
    src_name: Optional[str] = None,
    hide: bool = False,
    json_schema_extra: Optional[Dict[str, Any]] = None,
    **field_kwargs: Any,
):
    extra: Dict[str, Any] = dict(json_schema_extra or {})
    if target_name is not None:
        extra = _merge_json_extra(
            extra,
            {
                "direction": {
                    "target_name": target_name,
                    "src_name": src_name,
                    "targets": \
                        targets if isinstance(targets, tuple) else (targets,),
                }
            },
        )
    if hide:
        extra = _merge_json_extra(extra, {"hide": True})
    return PydanticField(default, json_schema_extra=extra, **field_kwargs)

T = TypeVar("T", bound=BaseModel)

class DirectFieldRule(BaseModel):
    attr_name: str
    src_name: str
    src_annotation: Type
    target_name: str
    targets: Tuple[Type[BaseModel], ...]

class MultiPydanticInheritBase(BaseModel):
    _direction_rules: Tuple[DirectFieldRule, ...] = ()
    _hidden_fields: Tuple[str, ...] = ()

    @classmethod
    def __pydantic_init_subclass__(cls, **kwargs):
        super().__pydantic_init_subclass__(**kwargs)
        logger.debug(
            "=" * 20 +\
            f"Start checking `{cls.__name__}` for multi-base model inherit" +\
            "=" * 20)
        # 1) Collect BaseModel parents except BaseModel itself
        parents: List[Type[BaseModel]] = [
            b
            for b in cls.__mro__[1:]
            if isinstance(b, type) and issubclass(b, BaseModel)
            and b is not BaseModel
        ]

        # 2) Collect per-field rules from json_schema_extra
        dir_rules: List[DirectFieldRule] = []
        from_rule_dict: Dict[Type[BaseModel], Dict[str, DirectFieldRule]] = {}
        hidden: List[str] = []
        for name, fld in getattr(cls, "model_fields", {}).items():
            extra = getattr(fld, "json_schema_extra", None) or {}
            if extra.get("hide"):
                hidden.append(name)
            d = extra.get("direction")
            if isinstance(d, dict) and "target_name" in d:
                rule = DirectFieldRule(
                    attr_name=name,
                    src_name=d.get("src_name") or name,
                    src_annotation=fld.annotation,
                    target_name=d["target_name"],
                    targets=tuple(d.get("targets") or ()),
                )
                dir_rules.append(rule)
                for t in rule.targets:
                    from_rule_dict.setdefault(t, {})[rule.target_name] = rule
                    logger.debug(
                        f"Direction `{rule.src_name}({cls.__name__})` ->"
                        f" `{rule.target_name}({t.__name__})` added")

        cls._direction_rules = tuple(dir_rules)
        cls._hidden_fields = tuple(hidden)

        errs: List[str] = []
        # 3) Type consistency check among parents
        seen: Dict[str, List[Tuple[Type[BaseModel], Any]]] = {}
        for p in parents:
            from_rule = from_rule_dict.get(p, {}) 
            for name, fld in getattr(p, "model_fields", {}).items():
                if name not in from_rule:
                    seen.setdefault(name, []).append((p, fld))
                    logger.debug(f"See the field `{name}` from"
                                   f" `{p.__name__}`")
                else:
                    dir_rule = from_rule[name]
                    seen.setdefault(dir_rule.src_name, []).append((p, fld))
                    logger.debug(
                        f"See the field `{name}(<- {dir_rule.src_name})`"
                        f" from `{p.__name__}`")
                    if dir_rule.src_annotation != fld.annotation:
                        errs.append(
                            f"Direct field '{dir_rule.src_name}' type mismatch:"
                            f" [{dir_rule.src_annotation!r} in"
                            f" {cls.__name__}] vs [{fld.annotation!r} in"
                            f" {p.__name__}]"
                        )

        if getattr(cls, "require_same_type", True):
            for name, items in seen.items():
                if len(items) < 2:
                    continue
                n0 = items[0][0].__name__
                t0 = items[0][1].annotation
                for p, fld in items[1:]:
                    n1 = p.__name__
                    t1 = fld.annotation
                    if fld.annotation != t0:
                        errs.append(
                            f"Same field '{name}' but type mismatch: "
                            f"[{t0!r} in {n0}] vs [{t1!r} in {n1}]"
                        )
        if errs:
            raise TypeError(
                f"{cls.__name__} conflicts:\n- " + "\n- ".join(errs)
            )


        # 4) Install hide getters/dump once per subclass
        if hidden:
            if not hasattr(cls, "__wrapped_getattribute__"):
                cls.__wrapped_getattribute__ = cls.__getattribute__  # type: ignore

                def __getattribute__(self, key: str):
                    if key in cls._hidden_fields:
                        err_msg = f"Field '{key}' is hidden on {cls.__name__}."
                        raise TypeError(err_msg)
                    return cls.__wrapped_getattribute__(self, key)  # type: ignore

                cls.__getattribute__ = __getattribute__  # type: ignore

            if not hasattr(cls, "__wrapped_model_dump__"):
                cls.__wrapped_model_dump__ = cls.model_dump  # type: ignore

                def model_dump(self, *a: Any, **kw: Any):
                    excl = set(kw.pop("exclude", set())) | set(cls._hidden_fields)
                    return cls.__wrapped_model_dump__(  # type: ignore
                        self, *a, exclude=excl, **kw
                    )

                cls.model_dump = model_dump  # type: ignore
            
            if not hasattr(cls, "__wrapped_setattr__"):
                cls.__wrapped_setattr__ = cls.__setattr__  # type: ignore

                def __setattr__(self, key: str, value: Any):
                    print(f"Setting attribute: {key}")
                    if key in cls._hidden_fields:
                        raise TypeError(
                            f"Field '{key}' is hidden on {cls.__name__}."
                        )
                    return cls.__wrapped_setattr__(self, key, value)  # type: ignore

                cls.__setattr__ = __setattr__  # type: ignore
                
        logger.debug(
            "=" * 20 +\
            f"End checking `{cls.__name__}` for multi-base model inherit" +\
            "=" * 20)

    def dispatch(self, target_cls: Type[T]) -> T:
        data = self.model_dump()
        keep = set(target_cls.model_fields.keys())
        view = {k: v for k, v in data.items() if k in keep}

        for r in type(self)._direction_rules:
            if target_cls in r.targets:
                src = r.src_name
                tgt = r.target_name
                if src in data:
                    view[tgt] = data[src]

        target = target_cls.model_construct(**view)

        target = self.post_dispath(target)
        target.model_validate(target.model_dump())

        return target

    def post_dispath(self, instance: T) -> T:
        return instance
