"""Configuration schema (pydantic) for one simulation run. Loaded from YAML; hashed into manifests."""
from __future__ import annotations

import hashlib
import json
from pathlib import Path
from typing import Any

import yaml
from pydantic import BaseModel, Field, model_validator


class AssetsCfg(BaseModel):
    n_assets: int = 1
    tick_cents: int = 1
    initial_price_cents: int = 20_000      # 200.00 per lot; chosen to equal mu/r (= 50/0.0025) by default
    price_band_pct: float = 0.5
    max_order_qty: int = 10_000


class FundamentalsCfg(BaseModel):
    div_interval: int = 20
    r_period: float = 0.0025
    mu: list[float] = [50.0]
    mu_shift: float = 0.0
    p_switch: float = 0.0
    rho: float = 0.98
    sigma: float = 2.0
    jump_prob: float = 0.0
    jump_sigma: float = 0.0
    factor_loading: list[float] = [0.0]
    x0: float = 0.0


class NewsCfg(BaseModel):
    rate: float = 0.02
    noise_sd: float = 1.0
    misleading_share: float = 0.0
    misleading_bias: float = 10.0
    window: int = 20


class ExchangeCfg(BaseModel):
    maker_fee_bps: float = 0.0
    taker_fee_bps: float = 2.0
    min_latency: float = 0.05
    halt_threshold: float | None = None
    halt_window: int = 10
    halt_duration: int = 20
    halt_cancel_resting: bool = True
    allow_short: bool = False
    max_short_lots: int = 0
    margin_check: bool = True
    liquidation_fraction: float = 0.5
    order_ttl: int | None = 50
    position_limit: int | None = None      # max |position| in lots per agent (safeguard treatment)
    call_auction_interval: int | None = None   # if set, batch orders and clear at a uniform price every k ticks (robustness)


class EndowmentCfg(BaseModel):
    cash_cents: int = 2_000_000           # 20,000.00
    shares: int = 100
    wealth_dispersion: float = 0.3        # lognormal sd of the wealth multiplier across traders


class LatencyCfg(BaseModel):
    min: float = 0.05
    max: float = 1.0
    llm_latency: float | None = None      # matched-latency design: LLM agents get this latency (None = drawn like others)


class LLMCfg(BaseModel):
    provider_mix: dict[str, float] = {"mock": 1.0}      # provider -> share of LLM slots
    models: dict[str, str] = {"mock": "mock-v1", "openai": "gpt-5.4-nano", "anthropic": "claude-haiku-5-5"}
    prompt_set: str = "single"            # 'single' | 'diverse'
    prompt_version: str = "v1"
    memory: int = 0                       # number of past own-decision summaries shown (0 = none)
    channel: bool = False                 # cohort communication channel
    cohort_size: int = 5
    decision_interval: int = 20
    stagger: bool = True                  # spread decision offsets across the interval
    stagger_slots: int = 4                # number of distinct offsets (batches of agents decide together -> parallel API calls)
    informed_share: float | None = None   # None: an LLM agent inherits the information set of the slot it replaces (default, keeps
                                          # information fixed within matched pairs); a number overrides it (information-equalization arm)
    precision: float = 2.0
    max_order_qty: int = 50
    temperature: float | None = None
    reasoning_effort: str | None = "low"  # openai gpt-5 family; anthropic 'effort'
    max_output_tokens: int = 1200
    concurrency: int = 8
    timeout_s: float = 60.0
    max_retries: int = 3
    usd_cap: float = 5.0                  # hard cap for this run (USD)
    token_cap: int | None = None
    decision_cap: int | None = None       # max decision opportunities per agent (fairness design)
    cache: bool = True                    # reuse identical (provider, model, prompt) decisions
    frozen_library: str | None = None     # path to a recorded decision library for offline replay
    adaptive_prompt: bool = False         # include P&L feedback line in the prompt
    price_assumptions: dict[str, dict[str, float]] = {
        # USD per 1M tokens: input, cached input, output — see docs/model_cards.md for sources and dates
        "gpt-5.4-nano": {"input": 0.20, "cached": 0.02, "output": 1.25},
        "gpt-5.4-mini": {"input": 0.75, "cached": 0.075, "output": 4.50},
        "gpt-5-nano": {"input": 0.05, "cached": 0.005, "output": 0.40},
        "gpt-5-mini": {"input": 0.25, "cached": 0.025, "output": 2.00},
        "gpt-4.1-nano": {"input": 0.10, "cached": 0.025, "output": 0.40},
        "gpt-4.1-mini": {"input": 0.40, "cached": 0.10, "output": 1.60},
        "claude-haiku-5-5": {"input": 0.10, "cached": 0.01, "output": 0.50},
        "claude-sonnet-5-5": {"input": 2.00, "cached": 0.20, "output": 10.00},
        "mock-v1": {"input": 0.0, "cached": 0.0, "output": 0.0},
    }


class PopulationCfg(BaseModel):
    n_traders: int = 100
    mix: dict[str, float] = {"fundamentalist": 0.25, "trend": 0.15, "meanrev": 0.10, "noise": 0.35,
                             "market_maker": 0.05, "statarb": 0.0, "leveraged": 0.10}
    activity: dict[str, float] = {"fundamentalist": 0.20, "trend": 0.20, "meanrev": 0.20, "noise": 0.30,
                                  "market_maker": 0.5, "statarb": 0.2, "leveraged": 0.2}
    policy_kwargs: dict[str, dict[str, Any]] = {}
    llm_share: float = 0.0                # fraction of trader *slots* controlled by LLM policies
    llm_replace: str = "proportional"     # 'proportional' (stratified across types) or a policy name
    llm_keep_market_makers: bool = True   # never replace market makers unless llm_replace == 'market_maker'
    informed_share: float = 0.3           # share of fundamentalists receiving private signals
    precision: float = 2.0
    leverage_cap: float = 3.0             # for 'leveraged' slots
    maintenance_margin: float = 0.25
    endowment: EndowmentCfg = EndowmentCfg()
    latency: LatencyCfg = LatencyCfg()
    llm: LLMCfg = LLMCfg()

    @model_validator(mode="after")
    def _check(self) -> "PopulationCfg":
        s = sum(self.mix.values())
        if abs(s - 1.0) > 1e-6:
            raise ValueError(f"population.mix must sum to 1 (got {s})")
        if not 0.0 <= self.llm_share <= 1.0:
            raise ValueError("llm_share must be in [0,1]")
        return self


class ShockCfg(BaseModel):
    time: int
    kind: str                     # fundamental_jump | rate | funding | liquidity | placebo | news_burst
    size: float = 0.0
    asset: int | None = None
    duration: int = 0


class OutputCfg(BaseModel):
    events: bool = False          # write the full event log (large)
    trades: bool = True
    llm_log: bool = True
    lookback: int = 60            # ticks of mid history visible to agents


class RunConfig(BaseModel):
    name: str = "run"
    seed: int = 1
    n_ticks: int = 500
    warmup: int = 50              # ticks excluded from outcome metrics
    assets: AssetsCfg = AssetsCfg()
    fundamentals: FundamentalsCfg = FundamentalsCfg()
    news: NewsCfg = NewsCfg()
    exchange: ExchangeCfg = ExchangeCfg()
    population: PopulationCfg = PopulationCfg()
    shocks: list[ShockCfg] = Field(default_factory=list)
    output: OutputCfg = OutputCfg()
    treatment: dict[str, Any] = Field(default_factory=dict)   # free-form labels copied to the manifest
    notes: str = ""

    # ---- helpers ---------------------------------------------------------------------------
    def config_hash(self) -> str:
        payload = json.dumps(self.model_dump(mode="json"), sort_keys=True).encode()
        return hashlib.sha256(payload).hexdigest()[:16]

    @classmethod
    def from_yaml(cls, path: str | Path, **overrides: Any) -> "RunConfig":
        with open(path) as f:
            raw = yaml.safe_load(f) or {}
        raw = deep_update(raw, overrides)
        return cls.model_validate(raw)

    def with_overrides(self, **overrides: Any) -> "RunConfig":
        raw = deep_update(self.model_dump(mode="json"), overrides)
        return RunConfig.model_validate(raw)

    def to_yaml(self, path: str | Path) -> None:
        with open(path, "w") as f:
            yaml.safe_dump(self.model_dump(mode="json"), f, sort_keys=False)


SECTIONS = {"assets", "fundamentals", "news", "exchange", "population", "llm", "endowment", "latency", "output", "treatment"}


def deep_update(base: dict, upd: dict) -> dict:
    """Recursive override. Config *sections* (nested models) are merged key by key; *leaf* dictionaries
    (provider_mix, mix, activity, models, policy_kwargs, price_assumptions, ...) are REPLACED wholesale, so
    `population.llm.provider_mix: {openai: 1.0}` yields exactly {openai: 1.0} and not a mock/openai mixture."""
    out = dict(base)
    for k, v in upd.items():
        if "." in k:   # dotted override: population.llm_share=0.25
            head, rest = k.split(".", 1)
            out[head] = deep_update(out.get(head, {}) or {}, {rest: v})
        elif isinstance(v, dict) and isinstance(out.get(k), dict) and k in SECTIONS:
            out[k] = deep_update(out[k], v)
        else:
            out[k] = v
    return out