from __future__ import annotations

import copy
import json
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.config import current_app_mode, is_answers_mode_on
from backend.constants import DEFAULT_EXTRACTION_MODE
from backend.models.admin_config import AdminConfig
from backend.utils.local_dev import is_local_dev
from backend.utils.local_dev_utils import load_local_config
from backend.utils.logging_utils import get_logger

DEFAULT_CONFIG: AdminConfig = {
    "version": 1,
    "required": [
        "agentHubLLM",
        "enterpriseAgentsDescription",
        "myAgentsTextCompletionModels",
        "myAgentsEmbeddingModel",
        "myAgentsFsConnection",
        "myAgentsNumDocs",
        "chartsMaxArtifactsSize",
        "uploadManagedFolder",
    ],
    "agentHubLLM": None,
    "agentHubOptionalInstructions": "",
    "orchestrationMode": "tools",
    "docGenerationEnabled": False,
    "smartMode": True,
    "permanentlyDeleteMessages": False,
    "dkuCallerTicket": True,
    "advancedSettingsEnabled": False,
    "includeUserAgentInContext": False,
    "includeClientIPInContext": False,
    "passOriginalRequestHeadersInContext": False,
    "allowDisableAgents": True,
    "logsLevel": "info",
    "conversationStarterExamples": [],
    "allowedEmbedOriginsEnabled": False,
    "allowedEmbedOrigins": [],
    "enableDocumentUpload": False,
    "extractionMode": DEFAULT_EXTRACTION_MODE,
    "imageValidation": True,
    "maxImagesInConversation": 50,
    "uploadManagedFolder": None,
    "conversationVisionLLM": None,
    "myAgentsEnabled": False,
    "myAgentsTextCompletionModels": [],
    "myAgentsEmbeddingModel": None,
    "myAgentsVisionLLM": None,
    "myAgentsFsConnection": None,
    "myAgentsSharepointConnection": None,
    "myAgentsSharepointSyncInterval": 1,
    "myAgentsSharepointSyncFrequency": "daily",
    "myAgentsAllowSharepointAgentSharing": False,
    "myAgentsFolder": None,
    "myAgentsNumDocs": 10,
    "myAgentsMaxUploadSizeMb": 10,
    "myAgentsManagedTools": [],
    "myAgentsEnablePromptLibrary": False,
    "myAgentsEnableChatWithTabularData": False,
    "myAgentsSemModelsInteractionMode": "simple",
    "myAgentsEnableDatasetSearchFromWorkspaces": True,
    "myAgentsEnableDatasetSearchFromCatalogs": True,
    "enterpriseAgents": [],
    "chartsGenerationMode": "auto",
    "chartsTextCompletionModel": None,
    "chartsAdditionalInstructions": None,
    "chartsMaxArtifactsSize": 2,
    "promptsAttributesById": {},
    "promptsAssetLibraryCollectionIds": ["DATAIKU"],
}


DEFAULT_CONFIG_ASW: AdminConfig = {
    "version": 1,
    "required": [
        "aswChatLLM",
        "chartsMaxArtifactsSize",
        "uploadManagedFolder",
    ],
    "agentHubLLM": None,
    "aswChatLLM": None,
    "agentHubOptionalInstructions": "",
    "docGenerationEnabled": False,
    "orchestratorLLM": None,
    "permanentlyDeleteMessages": False,
    "advancedSettingsEnabled": False,
    "includeUserAgentInContext": False,
    "includeClientIPInContext": False,
    "passOriginalRequestHeadersInContext": False,
    "logsLevel": "info",
    "conversationStarterExamples": [],
    "allowedEmbedOriginsEnabled": False,
    "allowedEmbedOrigins": [],
    "enableDocumentUpload": False,
    "extractionMode": DEFAULT_EXTRACTION_MODE,
    "imageValidation": True,
    "maxImagesInConversation": 50,
    "uploadManagedFolder": None,
    "conversationVisionLLM": None,
    "myAgentsVisionLLM": None,
    "chartsGenerationMode": "auto",
    "chartsTextCompletionModel": None,
    "chartsMaxArtifactsSize": 2,
}


logger = get_logger(__name__)


ADMIN_CONFIG_USER = "__GLOBAL__"

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

CONFIG_PATH: Optional[Path] = None


def _read_db_config() -> tuple[Optional[AdminConfig], bool]:
    try:
        from backend.database import crud

        settings = crud.get_admin_config(ADMIN_CONFIG_USER)
        if settings is None:
            return None, True
        return settings, True
    except Exception as e:
        logger.error("[read_config] database read failed (%s)", e)
        return None, False


def _write_db_config(config: dict) -> AdminConfig:
    from backend.database import crud

    return crud.upsert_admin_settings(ADMIN_CONFIG_USER, config)


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


def _local_config_is_old_config(d):
    # new webapp config only contains storage_type and log_level
    expected = {"storage_type", "log_level"}
    return not set(d.keys()).issubset(expected)


def get_default_config_for_mode() -> AdminConfig:
    if is_answers_mode_on():
        return DEFAULT_CONFIG_ASW
    return DEFAULT_CONFIG


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

    Args:
        old_config: Legacy configuration dictionary

    Returns:
        Typed AdminConfig dictionary
    """

    # ---------- 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["myAgentsSharepointConnection"] = None
    new_config["myAgentsSharepointSyncInterval"] = 1
    new_config["myAgentsSharepointSyncFrequency"] = "daily"
    new_config["myAgentsAllowSharepointAgentSharing"] = False
    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["chartsAdditionalInstructions"] = old_config.get("charts_prompt")
    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 setup_admin_config(workload_folder_path: Optional[str], app_mode: str = "agenthub") -> None:
    """
    Initialize admin config if it doesn't exist.
    """
    mode = str(app_mode or "agenthub").lower()
    logger.debug(f"[setup_admin_config] starting in mode {app_mode}")

    global _CONFIG_CACHE, CONFIG_PATH  # _CONFIG_ETAG, CONFIG_PATH
    CONFIG_FILENAME = "admin-config.json"
    config_path: Optional[Path] = None
    data = None
    migrated = None
    # Create config directory if workload_folder_path is provided, and set CONFIG_PATH for potential file-based migration.
    if workload_folder_path:
        config_dir = Path(workload_folder_path) / "config"
        CONFIG_PATH = config_dir / CONFIG_FILENAME
        config_path = CONFIG_PATH
        logger.info("[setup_admin_config] start (dir=%s, path=%s)", config_dir, CONFIG_PATH)
        config_dir.mkdir(parents=True, exist_ok=True)
    else:
        CONFIG_PATH = None
        logger.info("[setup_admin_config] start (no local workload folder; file migration disabled)")

    db_config, db_available = _read_db_config()
    # Case 1: DB has the config, no need for migration
    if db_available and db_config is not None:
        logger.info("[setup_admin_config] admin_settings found; loading")
        # Prepopulate cache with initial value
        _CONFIG_CACHE = db_config
        # _CONFIG_ETAG = _compute_etag(db_config)
        return

    # Case 2: Admin config file exists (from older versions) in the workload directory, migrate it to the database if available, and load it into cache
    if config_path and config_path.exists():
        logger.info("[setup_admin_config] admin-config.json found; migrating to admin_settings")
        raw = config_path.read_text(encoding="utf-8") or "{}"
        try:
            data = cast(AdminConfig, json.loads(raw))
        except json.JSONDecodeError:
            logger.error("[setup_admin_config] corrupted JSON; using default config for %s mode", mode)

        if db_available and isinstance(data, dict):
            try:
                written = _write_db_config(data)
                logger.info("[setup_admin_config] migration to admin_settings complete")
            except Exception as e:
                logger.warning("[setup_admin_config] failed to write admin_settings: %s", e)
                written = data
        else:
            logger.warning("[setup_admin_config] database unavailable; skipping admin_settings migration")
            written = data

        # Prepopulate cache with initial value
        _CONFIG_CACHE = written
        # _CONFIG_ETAG = _compute_etag(data)
        return

    # Case 3: No config in DB, no config file, first ever version where settings were defined in the UI of DSS
    if is_local_dev():
        local_config = load_local_config()
    else:
        try:
            local_config = get_webapp_config()
        except Exception as e:
            logger.warning("[setup_admin_config] could not get DSS webapp config: %s", e)
            local_config = None
    logger.info("[setup_admin_config] checking DSS webapp settings for migration")

    if local_config and isinstance(local_config, dict) and _local_config_is_old_config(local_config):
        logger.info("[setup_admin_config] migrating DSS webapp settings to admin_settings")
        new_config: AdminConfig = cast(AdminConfig, copy.deepcopy(get_default_config_for_mode()))
        migrated = migrate_webapp_config_to_admin_config(local_config, new_config)
    else:
        logger.info(
            "[setup_admin_config] no DSS webapp settings found; using default config for %s mode",
            mode,
        )

    if db_available and migrated is not None:
        try:
            migrated = _write_db_config(migrated)
            logger.info("[setup_admin_config] admin_settings write complete")
        except Exception as e:
            logger.warning("[setup_admin_config] failed to write admin_settings: %s", e)
    else:
        logger.warning("[setup_admin_config] database unavailable; skipping admin_settings write")

    _CONFIG_CACHE = migrated
    # _CONFIG_ETAG = _compute_etag(migrated)

    logger.info("[setup_admin_config] done")


def read_config() -> AdminConfig:
    """
    Read admin configuration from database (do not go through the cache, do not modify the cache)

    Returns:
        Tuple of (config_dict, etag)
    """
    logger.debug("[read_config] in (path=%s)", CONFIG_PATH)

    db_config, db_available = _read_db_config()

    if db_available and db_config is not None:
        data = db_config
    else:
        app_mode = current_app_mode()
        logger.warning(
            f"[read_config] {'admin_settings not found' if db_available else 'database unavailable'} ; returning default config for %s mode (not persisted)",
            app_mode,
        )
        data = get_default_config_for_mode()

    return data


def write_config(new_config: dict) -> None:
    """
    Write admin configuration to database.

    Args:
        new_config: Configuration dictionary to write

    Raises:
        ValueError: If config is not a dictionary
    """
    global _CONFIG_CACHE
    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")

    written_config = _write_db_config(new_config)
    # Force-Update in-memory cache.
    # This is an atomic swap of the variable, so no specific locking against get_cached_config is required
    _CONFIG_CACHE = written_config
    logger.info("[write_config] write OK (new_version=%s)", written_config.get("version"))


def get_cached_config() -> AdminConfig:
    """
    Get in-memory cached configuration. Falls back to read_config() if cache is empty and updates cache.
    """
    global _CONFIG_CACHE
    # A single variable access is atomic in Python, so this check is thread-safe without locking.
    # We accept that multiple threads might read from the database simultaneously on a cold cache.
    if _CONFIG_CACHE is not None:
        logger.debug("[get_cached_config] cache hit")
        return _CONFIG_CACHE

    data = read_config()
    _CONFIG_CACHE = data
    logger.info("[get_cached_config] cache miss; loaded config")
    return data
