#!/usr/bin/env python3
"""
Unified Strategy State — single authoritative source for all strategy management data.

Replaces the scattered multi-file strategy state (strategy_market_state.json,
active_instances.json, effective_params.json scheduler, runtime_scheduler_snapshot.json)
with ONE file: config/unified_strategy_state.json

Schema version 1:
  {
    "schema_version": 1,
    "strategies": { "<strategy_id>": { enabled, lifecycle_stage, is_custom, ... } },
    "instances": { "<instance_id>": { instance_id, strategy_id, status, allocation, ... } },
    "scheduler": { mode, active_strategy_ids, active_instance_ids, instance_weights, ... },
    "decision_policy": { profile, macro, interpreter, risk, execution, router, ... },
    "custom_workspace": { draft, generatedHistory, ... },
    "updated_at": "iso-timestamp"
  }

All reads go through load_unified_state() with a three-tier fallback:
  1. unified_strategy_state.json
  2. Legacy merge (strategy_market_state.json + active_instances.json + effective_params.json)
  3. Hardcoded default

All writes go through save_unified_state() using atomic_write_json.
"""

from __future__ import annotations

import json
import sys
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, Optional

from src.atomic_write import atomic_write_json

SCHEMA_VERSION = 1
UNIFIED_FILENAME = "unified_strategy_state.json"

# In-memory cache to avoid file I/O on every pipeline cycle
_cache: Dict[str, Any] = {}
_cache_ts: float = 0.0
_cache_ttl: float = 5.0  # seconds — short enough for GUI responsiveness, long enough for pipeline


def _unified_state_path(app_root: Optional[Path] = None) -> Path:
    if app_root is not None:
        return app_root / "config" / UNIFIED_FILENAME
    import config_manager
    return config_manager.get_config_dir() / UNIFIED_FILENAME


def _bundled_unified_state_path() -> Optional[Path]:
    if not getattr(sys, "frozen", False):
        return None
    meipass = getattr(sys, "_MEIPASS", "")
    if not meipass:
        return None
    path = Path(meipass) / "config" / UNIFIED_FILENAME
    return path if path.exists() else None


def _load_valid_state(path: Path) -> Optional[Dict[str, Any]]:
    try:
        data = json.loads(path.read_text(encoding="utf-8"))
    except Exception:
        return None
    if isinstance(data, dict) and data.get("schema_version") == SCHEMA_VERSION:
        return data
    return None


def _state_instance_count(data: Dict[str, Any]) -> int:
    instances = data.get("instances", {})
    return len(instances) if isinstance(instances, dict) else 0


def _enabled_active_instance_ids_by_strategy(app_root: Optional[Path] = None) -> Dict[str, set[str]]:
    import config_manager

    if app_root is not None:
        cfg = app_root / "config" if app_root.name.lower() != "config" else app_root
    else:
        cfg = config_manager.get_config_dir()
    path = cfg / "active_instances.json"
    if not path.exists():
        return {}
    try:
        payload = json.loads(path.read_text(encoding="utf-8"))
    except Exception:
        return {}
    raw_instances = payload.get("instances", []) if isinstance(payload.get("instances", []), list) else []
    by_strategy: Dict[str, set[str]] = {}
    for item in raw_instances:
        if not isinstance(item, dict) or not bool(item.get("enabled", True)):
            continue
        strategy_id = str(item.get("strategy_id", "")).strip()
        instance_id = str(item.get("id", "")).strip()
        if strategy_id and instance_id:
            by_strategy.setdefault(strategy_id, set()).add(instance_id)
    return by_strategy


def _sanitize_scheduler_identity(
    runtime_state: Dict[str, Any],
    *,
    drop_selected_identity: bool = False,
) -> Dict[str, Any]:
    scheduler = runtime_state.get("scheduler", {})
    if not isinstance(scheduler, dict):
        return runtime_state
    cleaned_scheduler = dict(scheduler)
    changed = False
    identity_keys = (
        "selected_instance_id",
        "selected_strategy_id",
        "entry_owner_instance_id",
        "entry_owner_strategy_id",
    )
    if drop_selected_identity:
        for key in identity_keys:
            if cleaned_scheduler.get(key, "") not in ("", None):
                changed = True
            cleaned_scheduler[key] = ""
    else:
        instances = runtime_state.get("instances", {})
        instances = instances if isinstance(instances, dict) else {}
        active_instance_ids = [
            str(item).strip()
            for item in cleaned_scheduler.get("active_instance_ids", [])
            if str(item).strip()
        ]

        def _normalize_identity(instance_key: str, strategy_key: str) -> None:
            nonlocal changed
            instance_id = str(cleaned_scheduler.get(instance_key, "") or "").strip()
            strategy_id = str(cleaned_scheduler.get(strategy_key, "") or "").strip()
            if not instance_id:
                if strategy_id:
                    cleaned_scheduler[strategy_key] = ""
                    changed = True
                return
            if active_instance_ids and instance_id not in active_instance_ids:
                cleaned_scheduler[instance_key] = ""
                cleaned_scheduler[strategy_key] = ""
                changed = True
                return
            row = instances.get(instance_id, {})
            actual_strategy_id = str(row.get("strategy_id", "") or "").strip() if isinstance(row, dict) else ""
            if actual_strategy_id and strategy_id != actual_strategy_id:
                cleaned_scheduler[strategy_key] = actual_strategy_id
                changed = True

        _normalize_identity("selected_instance_id", "selected_strategy_id")
        _normalize_identity("entry_owner_instance_id", "entry_owner_strategy_id")

    if not changed:
        return runtime_state
    cleaned = dict(runtime_state)
    cleaned["scheduler"] = cleaned_scheduler
    return cleaned


def _maybe_clean_mixed_runtime_instances(
    path: Path,
    runtime_state: Dict[str, Any],
    app_root: Optional[Path] = None,
) -> Dict[str, Any]:
    instances = runtime_state.get("instances", {})
    if not isinstance(instances, dict) or not instances:
        return runtime_state
    explicit_by_strategy = _enabled_active_instance_ids_by_strategy(app_root)
    if not explicit_by_strategy:
        return runtime_state
    dropped_ids: list[str] = []
    cleaned_instances: Dict[str, Any] = {}
    for instance_id, item in instances.items():
        iid = str(instance_id).strip()
        if not iid:
            continue
        row = item if isinstance(item, dict) else {}
        strategy_id = str(row.get("strategy_id", "")).strip()
        explicit_ids = explicit_by_strategy.get(strategy_id, set())
        if explicit_ids and iid not in explicit_ids and iid.startswith(f"inst_{strategy_id}"):
            dropped_ids.append(iid)
            continue
        cleaned_instances[iid] = row
    if not dropped_ids:
        return runtime_state
    scheduler = _build_scheduler_from_instances(cleaned_instances)
    selected_instance_id = next(iter(scheduler.get("active_instance_ids", [])), "")
    selected_strategy_id = ""
    if selected_instance_id:
        selected_row = cleaned_instances.get(selected_instance_id, {})
        if isinstance(selected_row, dict):
            selected_strategy_id = str(selected_row.get("strategy_id", "")).strip()
    scheduler["reason"] = "cleaned_mixed_runtime_instances"
    scheduler["selected_instance_id"] = selected_instance_id
    scheduler["selected_strategy_id"] = selected_strategy_id
    cleaned = dict(runtime_state)
    cleaned["instances"] = cleaned_instances
    cleaned["scheduler"] = scheduler
    cleaned["schema_version"] = SCHEMA_VERSION
    cleaned["updated_at"] = datetime.now(timezone.utc).isoformat()
    atomic_write_json(path, cleaned)
    return cleaned


def _looks_legacy_runtime_state(data: Dict[str, Any]) -> bool:
    instances = data.get("instances", {})
    if not isinstance(instances, dict) or not instances:
        return False
    if any(
        isinstance(inst, dict) and str(inst.get("config_revision", "")).strip().lower() == "active_instances_legacy"
        for inst in instances.values()
    ):
        return True
    instance_ids = [str(iid).strip() for iid in instances.keys() if str(iid).strip()]
    if instance_ids and all(iid.startswith("inst_") for iid in instance_ids):
        return True
    scheduler = data.get("scheduler", {})
    reason = str(scheduler.get("reason", "")).strip().lower() if isinstance(scheduler, dict) else ""
    return reason == "auto_built_from_instances"


def _merge_runtime_customizations(runtime_state: Dict[str, Any], bundled_state: Dict[str, Any]) -> Dict[str, Any]:
    merged = dict(bundled_state)
    bundled_strategies = dict(bundled_state.get("strategies", {})) if isinstance(bundled_state.get("strategies", {}), dict) else {}
    runtime_strategies = runtime_state.get("strategies", {}) if isinstance(runtime_state.get("strategies", {}), dict) else {}
    for strategy_id, strategy in runtime_strategies.items():
        if (
            isinstance(strategy, dict)
            and strategy.get("is_custom") is True
            and strategy_id not in bundled_strategies
        ):
            bundled_strategies[strategy_id] = strategy
    merged["strategies"] = bundled_strategies

    runtime_workspace = runtime_state.get("custom_workspace", {})
    if isinstance(runtime_workspace, dict) and runtime_workspace:
        merged["custom_workspace"] = _normalize_custom_workspace(runtime_workspace)

    runtime_policy = runtime_state.get("decision_policy", {})
    if isinstance(runtime_policy, dict) and runtime_policy:
        merged["decision_policy"] = runtime_policy

    return merged


def _should_take_over_runtime_state(runtime_state: Dict[str, Any], bundled_state: Dict[str, Any]) -> bool:
    if not _looks_legacy_runtime_state(runtime_state):
        return False
    bundled_count = _state_instance_count(bundled_state)
    runtime_count = _state_instance_count(runtime_state)
    if bundled_count > runtime_count:
        return True
    bundled_ids = set((bundled_state.get("instances", {}) or {}).keys()) if isinstance(bundled_state.get("instances", {}), dict) else set()
    runtime_ids = set((runtime_state.get("instances", {}) or {}).keys()) if isinstance(runtime_state.get("instances", {}), dict) else set()
    return bool(bundled_ids and bundled_ids != runtime_ids)


def _maybe_take_over_from_bundled(path: Path, runtime_state: Dict[str, Any]) -> Dict[str, Any]:
    bundled_path = _bundled_unified_state_path()
    if bundled_path is None:
        return runtime_state
    bundled_state = _load_valid_state(bundled_path)
    if bundled_state is None:
        return runtime_state
    if not _should_take_over_runtime_state(runtime_state, bundled_state):
        return runtime_state
    migrated = _merge_runtime_customizations(runtime_state, bundled_state)
    migrated["schema_version"] = SCHEMA_VERSION
    migrated["updated_at"] = datetime.now(timezone.utc).isoformat()
    atomic_write_json(path, migrated)
    return migrated


# ── Core load/save ─────────────────────────────────────────────────


def load_unified_state(app_root: Optional[Path] = None, *, use_cache: bool = True) -> Dict[str, Any]:
    """Primary entry point for ALL strategy state reads.

    Returns the unified dict, falling back through legacy sources if the unified
    file is missing or corrupt.
    """
    global _cache, _cache_ts

    if use_cache and _cache and (time.time() - _cache_ts) < _cache_ttl:
        return _cache

    path = _unified_state_path(app_root)

    # Tier 1: unified file
    if path.exists():
        try:
            data = json.loads(path.read_text(encoding="utf-8"))
            if isinstance(data, dict) and data.get("schema_version") == SCHEMA_VERSION:
                data = _maybe_take_over_from_bundled(path, data)
                data = _maybe_clean_mixed_runtime_instances(path, data, app_root)
                sanitized = _sanitize_scheduler_identity(data)
                if sanitized is not data:
                    data = sanitized
                    atomic_write_json(path, data)
                _cache = data
                _cache_ts = time.time()
                return data
        except Exception:
            pass  # fall through to legacy

    # Tier 2: legacy merge
    data = _legacy_merge(app_root)
    if data:
        _cache = data
        _cache_ts = time.time()
        return data

    # Tier 3: hardcoded default
    data = _default_state()
    _cache = data
    _cache_ts = time.time()
    return data


def save_unified_state(data: Dict[str, Any], app_root: Optional[Path] = None) -> bool:
    """Atomically write the unified state. Invalidates cache."""
    global _cache_ts
    path = _unified_state_path(app_root)
    data["schema_version"] = SCHEMA_VERSION
    data["updated_at"] = datetime.now(timezone.utc).isoformat()
    ok = atomic_write_json(path, data)
    if ok:
        _cache_ts = 0.0  # invalidate cache
    return ok


def invalidate_cache() -> None:
    """Force next load_unified_state() to re-read from disk."""
    global _cache_ts
    _cache_ts = 0.0


# ── Access helpers — single extraction point for each domain ───────


def get_strategies(state: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
    """Return strategies dict keyed by strategy_id."""
    s = state if state is not None else load_unified_state()
    strategies = s.get("strategies", {})
    return strategies if isinstance(strategies, dict) else {}


def get_instances(state: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
    """Return instances dict keyed by instance_id."""
    s = state if state is not None else load_unified_state()
    instances = s.get("instances", {})
    return instances if isinstance(instances, dict) else {}


def get_scheduler(state: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
    """Return scheduler block."""
    s = state if state is not None else load_unified_state()
    scheduler = s.get("scheduler", {})
    return scheduler if isinstance(scheduler, dict) else {}


def get_decision_policy(state: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
    """Return decision_policy block."""
    s = state if state is not None else load_unified_state()
    dp = s.get("decision_policy", {})
    return dp if isinstance(dp, dict) else {}


def get_custom_workspace(state: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
    """Return custom strategy workspace block."""
    s = state if state is not None else load_unified_state()
    workspace = s.get("custom_workspace", {})
    if not isinstance(workspace, dict):
        workspace = {}
    return _normalize_custom_workspace(workspace)


def get_running_strategy_ids(state: Optional[Dict[str, Any]] = None) -> list[str]:
    """Return strategy_ids of all running (non-paused) instances, deduplicated."""
    instances = get_instances(state)
    seen: set[str] = set()
    for inst in instances.values():
        if not isinstance(inst, dict):
            continue
        if str(inst.get("status", "")).lower() != "running":
            continue
        sid = str(inst.get("strategy_id", "")).strip()
        if sid:
            seen.add(sid)
    result = list(seen)
    return result


def get_running_instance_ids(state: Optional[Dict[str, Any]] = None) -> list[str]:
    """Return instance_ids of all running instances."""
    instances = get_instances(state)
    return [
        iid for iid, inst in instances.items()
        if isinstance(inst, dict) and str(inst.get("status", "")).lower() == "running"
    ]


def get_running_instances(state: Optional[Dict[str, Any]] = None) -> list[Dict[str, Any]]:
    """Return list of running instance dicts (list form for pipeline compat)."""
    instances = get_instances(state)
    return [
        dict(inst) for inst in instances.values()
        if isinstance(inst, dict) and str(inst.get("status", "")).lower() == "running"
    ]


def get_strategy(state: Optional[Dict[str, Any]] = None, strategy_id: str = "") -> Dict[str, Any]:
    """Return a single strategy entry by id, or empty dict."""
    strategies = get_strategies(state)
    entry = strategies.get(strategy_id, {})
    return entry if isinstance(entry, dict) else {}


def get_instance(state: Optional[Dict[str, Any]] = None, instance_id: str = "") -> Dict[str, Any]:
    """Return a single instance entry by id, or empty dict."""
    instances = get_instances(state)
    entry = instances.get(instance_id, {})
    return entry if isinstance(entry, dict) else {}


# ── Backward-compat: list-form access (for code that expects lists) ──


def instances_as_list(state: Optional[Dict[str, Any]] = None) -> list[Dict[str, Any]]:
    """Return all instances as a list (old runtime_instances format)."""
    instances = get_instances(state)
    return list(instances.values())


def strategies_as_list(state: Optional[Dict[str, Any]] = None) -> list[Dict[str, Any]]:
    """Return all strategies as a list with strategy_id included."""
    strategies = get_strategies(state)
    return [{"strategy_id": sid, **spec} for sid, spec in strategies.items()]


# ── Legacy fallback ────────────────────────────────────────────────


def _legacy_merge(app_root: Optional[Path] = None) -> Dict[str, Any]:
    """[DEPRECATED] Merge from legacy files when unified file is missing.

    This function fires when unified_strategy_state.json is unavailable.
    After successful merge, it auto-saves to unified format so subsequent
    calls use the canonical source.
    """
    import config_manager

    try:
        from bridge.gold_trader import _daemon_log
        _daemon_log("[LEGACY_MERGE] _legacy_merge invoked — unified_strategy_state.json is missing or corrupt", "WARN")
    except Exception:
        pass

    if app_root is not None:
        cfg = app_root / "config" if app_root.name.lower() != "config" else app_root
    else:
        cfg = config_manager.get_config_dir()

    # Try strategy_market_state.json as base
    market_path = cfg / "strategy_market_state.json"
    if not market_path.exists():
        return {}

    try:
        market = json.loads(market_path.read_text(encoding="utf-8"))
    except Exception:
        return {}

    if not isinstance(market, dict):
        return {}

    strategies = market.get("strategies", {}) if isinstance(market.get("strategies", {}), dict) else {}
    runtime_list = market.get("runtime_instances", []) if isinstance(market.get("runtime_instances", []), list) else []
    decision_policy = market.get("decision_policy", {}) if isinstance(market.get("decision_policy", {}), dict) else {}
    custom_workspace = _load_legacy_custom_workspace(cfg)

    # Track how many times legacy merge has been invoked
    _legacy_merge_count = int(market.get("_legacy_merge_count", 0) or 0) + 1

    # Convert runtime_instances list to dict
    instances: Dict[str, Any] = {}
    for inst in runtime_list:
        if not isinstance(inst, dict):
            continue
        iid = str(inst.get("instance_id", "")).strip()
        if iid:
            instances[iid] = inst

    # Merge from active_instances.json
    active_path = cfg / "active_instances.json"
    if active_path.exists():
        try:
            active = json.loads(active_path.read_text(encoding="utf-8"))
        except Exception:
            active = {}
        if isinstance(active, dict):
            for inst in active.get("instances", []) or []:
                if not isinstance(inst, dict):
                    continue
                iid = str(inst.get("id", "")).strip()
                sid = str(inst.get("strategy_id", "")).strip()
                if iid and sid and iid not in instances:
                    instances[iid] = {
                        "instance_id": iid,
                        "strategy_id": sid,
                        "strategy_name": sid,
                        "allocation": float(inst.get("allocation", 1.0) or 1.0),
                        "status": "running" if inst.get("enabled", True) else "paused",
                        "role": "live_support",
                    }

    # Merge scheduler from effective_params.json
    scheduler: Dict[str, Any] = {}
    logs_eff = cfg.parent / "logs" / "effective_params.json"
    if not logs_eff.exists():
        logs_eff = (
            ((app_root / "logs") if app_root and app_root.name.lower() != "logs" else app_root) / "effective_params.json"
            if app_root
            else None
        )
    if logs_eff and logs_eff.exists():
        try:
            eff = json.loads(logs_eff.read_text(encoding="utf-8"))
        except Exception:
            eff = {}
        if isinstance(eff, dict):
            s = eff.get("runtime_scheduler", {})
            if isinstance(s, dict):
                scheduler = dict(s)

    if not scheduler:
        # Build scheduler from instances
        scheduler = _build_scheduler_from_instances(instances)

    return _sanitize_scheduler_identity({
        "schema_version": SCHEMA_VERSION,
        "strategies": strategies,
        "instances": instances,
        "scheduler": scheduler,
        "decision_policy": decision_policy,
        "custom_workspace": custom_workspace,
        "updated_at": datetime.now(timezone.utc).isoformat(),
        "_legacy_merge_count": _legacy_merge_count,
    }, drop_selected_identity=True)


def _build_scheduler_from_instances(instances: Dict[str, Any]) -> Dict[str, Any]:
    """Build a minimal scheduler block from running instances."""
    active_iids: list[str] = []
    active_sids: list[str] = []
    weights: Dict[str, float] = {}
    sid_map: Dict[str, list[str]] = {}

    for iid, inst in instances.items():
        if not isinstance(inst, dict):
            continue
        if str(inst.get("status", "")).lower() != "running":
            continue
        sid = str(inst.get("strategy_id", "")).strip()
        if not sid:
            continue
        active_iids.append(iid)
        if sid not in active_sids:
            active_sids.append(sid)
        weights[iid] = float(inst.get("allocation", 1.0) or 1.0)
        sid_map.setdefault(sid, []).append(iid)

    return {
        "mode": "multi_active" if len(active_sids) > 1 else "single_active",
        "arbiter": "allocation-weighted",
        "reason": "auto_built_from_instances",
        "active_instance_ids": active_iids,
        "active_strategy_ids": active_sids,
        "instance_weights": weights,
        "strategy_weights": {sid: sum(weights.get(iid, 1.0) for iid in iids) for sid, iids in sid_map.items()},
        "strategy_instance_map": sid_map,
        "instances": [inst for inst in instances.values() if isinstance(inst, dict)],
    }


def _default_state() -> Dict[str, Any]:
    """Hardcoded default when all sources are unavailable."""
    return {
        "schema_version": SCHEMA_VERSION,
        "strategies": {
            "ema_macd": {
                "enabled": True,
                "lifecycle_stage": "paper",
                "is_custom": False,
                "display_name": "EMA+MACD+RSI+BB",
            },
        },
        "instances": {},
        "scheduler": {
            "mode": "idle",
            "arbiter": "allocation-weighted",
            "reason": "hardcoded_default",
            "active_instance_ids": [],
            "active_strategy_ids": [],
            "instance_weights": {},
            "strategy_weights": {},
            "strategy_instance_map": {},
        },
        "decision_policy": {},
        "custom_workspace": _normalize_custom_workspace({}),
        "updated_at": datetime.now(timezone.utc).isoformat(),
    }


def _normalize_custom_workspace(payload: Dict[str, Any]) -> Dict[str, Any]:
    workspace = payload if isinstance(payload, dict) else {}
    draft = workspace.get("draft", {}) if isinstance(workspace.get("draft", {}), dict) else {}
    normalized = {
        "draft": {
            "name": str(draft.get("name", "") or ""),
            "document": str(draft.get("document", "") or ""),
            "compileTarget": str(draft.get("compileTarget", "template") or "template"),
            "strategySpec": draft.get("strategySpec", {}) if isinstance(draft.get("strategySpec", {}), dict) else {},
            "optimized": draft.get("optimized", {}) if isinstance(draft.get("optimized", {}), dict) else {},
            "evaluation": draft.get("evaluation", {}) if isinstance(draft.get("evaluation", {}), dict) else {},
            "validationReport": draft.get("validationReport", {}) if isinstance(draft.get("validationReport", {}), dict) else {},
            "workflowMeta": draft.get("workflowMeta", {}) if isinstance(draft.get("workflowMeta", {}), dict) else {},
            "saveState": str(draft.get("saveState", "saved") or "saved"),
            "lastSavedAt": str(draft.get("lastSavedAt", "") or ""),
        },
        "generatedHistory": workspace.get("generatedHistory", []) if isinstance(workspace.get("generatedHistory", []), list) else [],
    }
    for key in ("conversation", "knowledge_state", "mode", "name", "document"):
        value = workspace.get(key)
        if key in {"conversation", "knowledge_state"}:
            if isinstance(value, dict) and value:
                normalized[key] = value
        elif isinstance(value, str) and value:
            normalized[key] = value
    return normalized


def _load_legacy_custom_workspace(cfg: Path) -> Dict[str, Any]:
    custom_path = cfg / "strategy_custom_state.json"
    if not custom_path.exists():
        return _normalize_custom_workspace({})
    try:
        data = json.loads(custom_path.read_text(encoding="utf-8"))
    except Exception:
        return _normalize_custom_workspace({})
    return _normalize_custom_workspace(data if isinstance(data, dict) else {})
