from __future__ import annotations

import copy
import json
import os
import tempfile
import threading
import time
from contextlib import contextmanager
from hashlib import sha256
from pathlib import Path
from typing import Optional, Tuple, cast

from dataiku.customwebapp import get_webapp_config

from backend.app_paths import get_workload_folder_path, is_local_dev
from backend.constants import DEFAULT_EXTRACTION_MODE
from backend.models.admin_config import AdminConfig
from backend.utils.local_dev_utils import load_local_config
from backend.utils.logging_utils import get_logger

CONFIG_FILENAME = "admin-config.json"
CONFIG_DIR = Path(get_workload_folder_path()) / "config"
CONFIG_PATH = CONFIG_DIR / CONFIG_FILENAME

DEFAULT_CONFIG: AdminConfig = {
    "version": "1.0.0",
    "required": [
        "agentHubLLM",
        "enterpriseAgentsDescription",
        "myAgentsTextCompletionModels",
        "myAgentsEmbeddingModel",
        "myAgentsFsConnection",
        "myAgentsNumDocs",
        "chartsMaxArtifactsSize",
        "uploadManagedFolder",
    ],
    "agentHubLLM": None,
    "agentHubOptionalInstructions": "",
    "orchestrationMode": "tools",
    "smartMode": True,
    "permanentlyDeleteMessages": False,
    "allowDisableAgents": True,
    "logsLevel": "info",
    "conversationStarterExamples": [],
    "enableDocumentUpload": False,
    "extractionMode": DEFAULT_EXTRACTION_MODE,
    "maxImagesInConversation": 50,
    "uploadManagedFolder": None,
    "conversationVisionLLM": None,
    "myAgentsEnabled": False,
    "myAgentsTextCompletionModels": [],
    "myAgentsEmbeddingModel": None,
    "myAgentsFsConnection": None,
    "myAgentsFolder": None,
    "myAgentsNumDocs": 10,
    "myAgentsManagedTools": [],
    "myAgentsEnablePromptLibrary": False,
    "enterpriseAgents": [],
    "chartsGenerationMode": "auto",
    "chartsTextCompletionModel": None,
    "chartsMaxArtifactsSize": 2,
    "promptsAttributesById": {},
}

logger = get_logger(__name__)

_LOCK = threading.RLock()

# In-memory cache
_CONFIG_CACHE: Optional[AdminConfig] = None
_CONFIG_ETAG: Optional[str] = None


def _canonical_json(obj: dict) -> str:
    """Serialize dict to canonical JSON string."""
    return json.dumps(obj, separators=(",", ":"), sort_keys=True, ensure_ascii=False)


def _compute_etag(obj: dict) -> str:
    """Compute SHA256 hash of canonical JSON."""
    return sha256(_canonical_json(obj).encode("utf-8")).hexdigest()


@contextmanager
def _with_lock():
    """Serialize access to the config file within this process."""
    tname = threading.current_thread().name
    logger.debug("[lock] waiting (thread=%s)", tname)
    start = time.time()
    with _LOCK:
        waited = time.time() - start
        if waited >= 0.005:
            logger.debug("[lock] acquired after %.3fs (thread=%s)", waited, tname)
        else:
            logger.debug("[lock] acquired (thread=%s)", tname)
        try:
            yield
        finally:
            logger.debug("[lock] released (thread=%s)", tname)


def migrate_webapp_config_to_admin_config(old_config: dict) -> AdminConfig:
    """
    Migrate legacy webapp.json config to new flat admin-config.json format.

    Args:
        old_config: Legacy configuration dictionary

    Returns:
        Typed AdminConfig dictionary
    """
    new_config: AdminConfig = cast(AdminConfig, copy.deepcopy(DEFAULT_CONFIG))

    # ---------- LLMs / My Agents Text Completion Models ----------

    text_ids = []
    for entry in old_config.get("LLMs", []) or []:
        if isinstance(entry, dict) and entry.get("llm_id"):
            text_ids.append(entry["llm_id"])
        elif isinstance(entry, str):
            text_ids.append(entry)

    id_to_name = {}
    if text_ids:
        try:
            import dataiku

            proj = dataiku.api_client().get_default_project()
            for llm in proj.list_llms("GENERIC_COMPLETION"):
                lid = llm.get("id")
                if lid:
                    id_to_name[lid] = llm.get("friendlyName") or lid
        except Exception as e:
            logger.warning("[migration] could not resolve friendly names from DSS: %s", e)

        new_config["myAgentsTextCompletionModels"] = [{"id": lid, "name": id_to_name.get(lid, lid)} for lid in text_ids]

    # Embedding model
    new_config["myAgentsEmbeddingModel"] = old_config.get("embedding_llm")

    # Agent Hub LLM settings
    new_config["agentHubLLM"] = old_config.get("default_llm_id")
    new_config["agentHubOptionalInstructions"] = old_config.get("globalSystemPrompt", "")
    new_config["orchestrationMode"] = old_config.get("orchestration_mode", "tools")

    # ---------- Enterprise Agents ----------
    items = []

    # Tool agents
    for cfg in old_config.get("tool_agent_configurations", []) or []:
        if not isinstance(cfg, dict):
            continue
        agent_id = cfg.get("agent_id")
        if not agent_id:
            continue
        toks = agent_id.split(":")
        items.append(
            {
                "id": toks[2],
                "projectKey": toks[0],
                "type": toks[1],
                "name": cfg.get("tool_agent_display_name", ""),
                "description": cfg.get("tool_agent_description", ""),
                "additionalInstructions": cfg.get("agent_system_instructions", ""),
                "exampleQuestions": cfg.get("agent_example_queries", []),
                "allowInsightsCreation": bool(cfg.get("enable_stories", False)),
                "storiesWorkspace": cfg.get("stories_workspace") if cfg.get("enable_stories") else None,
            }
        )

    # Augmented LLMs
    for cfg in old_config.get("augmented_llms_configurations", []) or []:
        if not isinstance(cfg, dict):
            continue
        augmented_llms_id = cfg.get("augmented_llm_id")
        if not augmented_llms_id:
            continue
        toks = augmented_llms_id.split(":")
        items.append(
            {
                "id": toks[2],
                "projectKey": toks[0],
                "type": toks[1],
                "name": cfg.get("augmented_llm_display_name", ""),
                "description": cfg.get("augmented_llm_description", ""),
                "additionalInstructions": cfg.get("augmented_llm_system_instructions", ""),
                "exampleQuestions": cfg.get("augllm_example_queries", []),
                "allowInsightsCreation": bool(cfg.get("enable_stories", False)),
                "storiesWorkspace": None,
            }
        )

    # Flat structure: enterpriseAgents is directly a list
    new_config["enterpriseAgents"] = [i for i in items if i.get("id")]

    # ---------- My Agents (User Agents) ----------
    new_config["myAgentsEnabled"] = bool(old_config.get("enable_quick_agents", False))
    new_config["myAgentsEnablePromptLibrary"] = bool(old_config.get("enable_prompt_library", False))
    new_config["myAgentsFsConnection"] = old_config.get("default_fs_connection")
    new_config["myAgentsFolder"] = old_config.get("agents_folder")
    new_config["myAgentsNumDocs"] = old_config.get("number_of_documents_to_retrieve")
    new_config["myAgentsManagedTools"] = old_config.get("tools", []) or []

    # ---------- Charts / Visualization ----------
    viz_mode = (old_config.get("visualization_generation_mode") or "AUTO").lower()
    new_config["chartsGenerationMode"] = viz_mode  # type: ignore
    new_config["chartsTextCompletionModel"] = old_config.get("charts_generation_llm_id") or old_config.get(
        "graph_generation_llm_id"
    )
    new_config["chartsMaxArtifactsSize"] = old_config.get("max_artifacts_size_mb", 2)

    # ---------- General Settings ----------
    new_config["smartMode"] = bool(old_config.get("enable_app_smart_mode", False))
    new_config["permanentlyDeleteMessages"] = bool(old_config.get("permanent_delete_messages", False))
    new_config["allowDisableAgents"] = bool(old_config.get("allow_chat_to_llm", True))

    try:
        new_config["maxImagesInConversation"] = int(old_config.get("quota_images_per_conversation", 50))
    except Exception:
        new_config["maxImagesInConversation"] = 50

    log_level = (old_config.get("logLevel") or "INFO").lower()
    new_config["logsLevel"] = log_level  # type: ignore

    # Uploads managed folder
    upload_folder = old_config.get("uploads")
    if upload_folder:
        new_config["uploadManagedFolder"] = upload_folder

    # Homepage examples (conversation starters)
    examples = []
    for ex in old_config.get("homepage_examples", []) or []:
        if not isinstance(ex, dict):
            continue
        agent_ids = []
        for agent_ref in ex.get("selected_agents", []) or []:
            if isinstance(agent_ref, str) and ":" in agent_ref:
                toks = agent_ref.split(":")
                if len(toks) >= 3:
                    agent_ids.append(toks[2])
            elif isinstance(agent_ref, str):
                agent_ids.append(agent_ref)

        examples.append(
            {
                "label": ex.get("homepage_example_label", ""),
                "query": ex.get("homepage_example_query", ""),
                "enterpriseAgents": agent_ids,
            }
        )
    new_config["conversationStarterExamples"] = examples

    # ---------- Prompts ----------
    new_config["promptsAttributesById"] = {}

    return new_config


def _atomic_write_json(path: Path, data: dict) -> None:
    """Write JSON atomically to avoid partial writes."""
    serialized = _canonical_json(data)
    dir_ = path.parent
    dir_.mkdir(parents=True, exist_ok=True)
    fd, tmp_path = tempfile.mkstemp(dir=str(dir_), prefix=path.name, suffix=".tmp")
    logger.debug("[write] temp (tmp=%s, size=%d bytes)", tmp_path, len(serialized))
    try:
        with os.fdopen(fd, "w", encoding="utf-8") as f:
            f.write(serialized)
        os.replace(tmp_path, path)
        logger.debug("[write] replace OK (target=%s)", path)
    finally:
        try:
            if os.path.exists(tmp_path):
                os.remove(tmp_path)
                logger.debug("[write] temp cleanup OK (tmp=%s)", tmp_path)
        except Exception as e:  # pragma: no cover
            logger.warning("[write] temp cleanup failed (%s): %s", tmp_path, e)


def setup_admin_config() -> None:
    """
    Initialize admin config if it doesn't exist.
    """
    global _CONFIG_CACHE, _CONFIG_ETAG
    logger.info("[setup_admin_config] start (dir=%s, path=%s)", CONFIG_DIR, CONFIG_PATH)
    CONFIG_DIR.mkdir(parents=True, exist_ok=True)

    with _with_lock():
        if CONFIG_PATH.exists():
            logger.info("[setup_admin_config] admin-config.json already exists; loading")
            # Load existing config into cache
            data, etag = read_config()
            _CONFIG_CACHE = data
            _CONFIG_ETAG = etag
            return

        if (is_local_dev()):
            old_config = load_local_config()
        else:
            old_config = get_webapp_config()
        logger.info("[setup_admin_config] checking DSS webapp settings for migration")

        if old_config and isinstance(old_config, dict):
            logger.info("[setup_admin_config] migrating DSS webapp settings to admin-config.json")
            migrated = migrate_webapp_config_to_admin_config(old_config)
            _atomic_write_json(CONFIG_PATH, migrated)
            # seed in-memory cache
            _CONFIG_CACHE = migrated
            _CONFIG_ETAG = _compute_etag(migrated)
            logger.info("[setup_admin_config] migration complete")
        else:
            logger.info("[setup_admin_config] no DSS webapp settings found; writing DEFAULT_CONFIG")
            _atomic_write_json(CONFIG_PATH, DEFAULT_CONFIG)
            # seed in-memory cache
            _CONFIG_CACHE = DEFAULT_CONFIG
            _CONFIG_ETAG = _compute_etag(DEFAULT_CONFIG)

    logger.info("[setup_admin_config] done")


def read_config() -> Tuple[AdminConfig, str]:
    """
    Read admin configuration from disk.

    Returns:
        Tuple of (config_dict, etag)
    """
    logger.debug("[read_config] in (path=%s)", CONFIG_PATH)
    with _with_lock():
        if not CONFIG_PATH.exists():
            logger.warning("[read_config] file not found; returning DEFAULT_CONFIG (not persisted)")
            data = DEFAULT_CONFIG
        else:
            raw = CONFIG_PATH.read_text(encoding="utf-8") or "{}"
            try:
                data = cast(AdminConfig, json.loads(raw))
            except json.JSONDecodeError:
                logger.error("[read_config] corrupted JSON; returning DEFAULT_CONFIG (not persisted)")
                data = DEFAULT_CONFIG

        etag = _compute_etag(data)
        logger.debug("[read_config] etag=%s", etag)
        return data, etag


def write_config(new_config: dict, *, expected_etag: Optional[str]) -> str:
    """
    Write admin configuration to disk with optimistic locking.

    Args:
        new_config: Configuration dictionary to write
        expected_etag: Expected etag for optimistic locking (None to skip check)

    Returns:
        New etag of written configuration

    Raises:
        ValueError: If config is not a dictionary
        PreconditionFailed: If etag mismatch
    """
    global _CONFIG_CACHE, _CONFIG_ETAG
    logger.debug("[write_config] in (expected_etag=%s)", expected_etag)
    if not isinstance(new_config, dict):
        logger.error("[write_config] payload must be a dict, got %s", type(new_config))
        raise ValueError("config must be an object")

    with _with_lock():
        current, current_etag = read_config()
        logger.debug("[write_config] current_etag=%s", current_etag)

        if expected_etag and expected_etag != current_etag:
            from werkzeug.exceptions import PreconditionFailed

            logger.warning("[write_config] etag mismatch (expected=%s != current=%s)", expected_etag, current_etag)
            raise PreconditionFailed("etag mismatch")

        _atomic_write_json(CONFIG_PATH, new_config)
        new_etag = _compute_etag(new_config)
        # Update in-memory cache
        _CONFIG_CACHE = cast(AdminConfig, new_config)
        _CONFIG_ETAG = new_etag
        logger.info("[write_config] write OK (new_etag=%s)", new_etag)
        return new_etag


def get_cached_config() -> Tuple[AdminConfig, str]:
    """
    Get in-memory cached configuration.
    Falls back to read_config() if cache is empty.

    Returns:
        Tuple of (config_dict, etag)
    """
    global _CONFIG_CACHE, _CONFIG_ETAG
    with _with_lock():
        if _CONFIG_CACHE is not None and _CONFIG_ETAG is not None:
            return _CONFIG_CACHE, _CONFIG_ETAG
        data, etag = read_config()
        _CONFIG_CACHE, _CONFIG_ETAG = data, etag
        return data, etag
