from dataclasses import dataclass, field
from typing import Dict, List, Optional
from enum import Enum


class ModelTier(Enum):
    SMALL = "small"
    HYBRID = "hybrid"
    POWERHOUSE = "powerhouse"


class ThinkingMode(Enum):
    DISABLED = "disabled"
    ENABLED = "enabled"
    AUTO = "auto"


@dataclass
class ModelConfig:
    role: str
    model_id: str
    fallback_model_id: Optional[str] = None
    max_cost_usd: float = 1.0
    timeout_ms: int = 30000
    max_retries: int = 2
    tier: ModelTier = ModelTier.HYBRID
    thinking_mode: ThinkingMode = ThinkingMode.DISABLED
    provider: str = "zai"


@dataclass
class Verdict:
    decision: str
    confidence: float
    rationale_tags: List[str] = field(default_factory=list)
    source_role: str = ""


class GateType(Enum):
    TOOL_CALL = "tool_call"
    MEMORY_WRITE = "memory_write"
    ESCALATION = "escalation"


@dataclass
class GateDecision:
    gate: GateType
    allowed: bool
    verdict: Verdict
    override_reason: Optional[str] = None


MODEL_POOL: Dict[str, ModelConfig] = {
    "router_small": ModelConfig(
        role="router_small",
        model_id="glm-4.7-flash",
        fallback_model_id=None,
        max_cost_usd=0.0,
        timeout_ms=10000,
        tier=ModelTier.SMALL,
        thinking_mode=ThinkingMode.DISABLED,
        provider="zai",
    ),
    "router_hybrid": ModelConfig(
        role="router_hybrid",
        model_id="glm-4.7-flash",
        fallback_model_id=None,
        max_cost_usd=0.0,
        timeout_ms=15000,
        tier=ModelTier.HYBRID,
        thinking_mode=ThinkingMode.ENABLED,
        provider="zai",
    ),
    "router_powerhouse": ModelConfig(
        role="router_powerhouse",
        model_id="glm-4.7-flash",
        fallback_model_id=None,
        max_cost_usd=0.0,
        timeout_ms=20000,
        tier=ModelTier.POWERHOUSE,
        thinking_mode=ThinkingMode.ENABLED,
        provider="zai",
    ),
    "verifier_v1": ModelConfig(
        role="verifier_v1",
        model_id="glm-4.7-flash",
        max_cost_usd=0.0,
        timeout_ms=10000,
        tier=ModelTier.SMALL,
        thinking_mode=ThinkingMode.DISABLED,
        provider="zai",
    ),
    "verifier_v2": ModelConfig(
        role="verifier_v2",
        model_id="glm-4.7-flash",
        max_cost_usd=0.0,
        timeout_ms=10000,
        tier=ModelTier.SMALL,
        thinking_mode=ThinkingMode.DISABLED,
        provider="zai",
    ),
    "verifier_v3": ModelConfig(
        role="verifier_v3",
        model_id="glm-4.7-flash",
        max_cost_usd=0.0,
        timeout_ms=10000,
        tier=ModelTier.SMALL,
        thinking_mode=ThinkingMode.DISABLED,
        provider="zai",
    ),
    "verifier_v4": ModelConfig(
        role="verifier_v4",
        model_id="glm-4.7-flash",
        max_cost_usd=0.0,
        timeout_ms=15000,
        tier=ModelTier.HYBRID,
        thinking_mode=ThinkingMode.ENABLED,
        provider="zai",
    ),
    "duplex_d1": ModelConfig(
        role="duplex_d1",
        model_id="glm-4.7-flash",
        max_cost_usd=0.0,
        timeout_ms=10000,
        tier=ModelTier.SMALL,
        thinking_mode=ThinkingMode.DISABLED,
        provider="zai",
    ),
    "duplex_d2": ModelConfig(
        role="duplex_d2",
        model_id="glm-4.7-flash",
        max_cost_usd=0.0,
        timeout_ms=10000,
        tier=ModelTier.SMALL,
        thinking_mode=ThinkingMode.DISABLED,
        provider="zai",
    ),
}

ALIAS_MAP: Dict[str, str] = {
    "small": "router_small",
    "hybrid": "router_hybrid",
    "powerhouse": "router_powerhouse",
}


def get_config(role: str) -> Optional[ModelConfig]:
    key = ALIAS_MAP.get(role, role)
    return MODEL_POOL.get(key)


def get_models_by_tier(tier: ModelTier) -> List[ModelConfig]:
    return [c for c in MODEL_POOL.values() if c.tier == tier]


def get_thinking_payload(mode: ThinkingMode) -> Optional[Dict[str, str]]:
    if mode == ThinkingMode.DISABLED:
        return {"type": "disabled"}
    elif mode == ThinkingMode.ENABLED:
        return {"type": "enabled"}
    return None


SLOT_SUMMARY = """
Model Pool Summary (TESTING MODE - all glm-4.7-flash):
+-- Router (3 tiers, 1 active per decision)
|   +-- router_small      -> glm-4.7-flash (no thinking)
|   +-- router_hybrid     -> glm-4.7-flash (thinking ON)
|   +-- router_powerhouse -> glm-4.7-flash (thinking ON)
|
+-- 4-way Consensus (4 parallel, outputs: Verdict)
|   +-- verifier_v1-v3    -> glm-4.7-flash (no thinking)
|   +-- verifier_v4       -> glm-4.7-flash (thinking ON)
|
+-- Duplex Verifier (2 roles, cross-check)
    +-- duplex_d1, d2     -> glm-4.7-flash (no thinking)

Gate Types:
+-- tool_call     -> Allow/deny tool execution
+-- memory_write  -> Allow/deny write to semantic memory / whiteboard
+-- escalation    -> Allow/deny Hybrid -> Powerhouse escalation

Field Sim:
+-- Deterministic (no model), callable by Scheduler/Hybrid for gradient guidance
"""
