"""DataSourceConfig API for user-defined data sources"""

from typing import Any, Optional
from datetime import datetime
import base64
import json
import re
from fastapi import APIRouter, Depends, HTTPException, Query, Request, status
from sqlalchemy import delete, select, func
from sqlalchemy.ext.asyncio import AsyncSession
from pydantic import BaseModel, Field
import httpx

from app.core.target_schema_registry import get_target_schema, list_target_schemas
from app.core.datasource_defaults import DEFAULT_DATASOURCES
from app.db.session import get_db
from app.models.user import User
from app.models.datasource_config import DataSourceConfig
from app.models.datasource_mapping import DataSourceMappingTemplate
from app.models.collected_data import CollectedData
from app.models.vessel import AISRawObservation, AISSourceHealth
from app.core.security import get_current_user
from app.core.cache import cache
from app.core.time import to_iso8601_utc
from app.ai_tasks.prompts import get_effective_prompt
from app.schemas.ai import SituationalAnalysisRequest
from app.services.ai_client import AIProviderClient, get_ai_provider_client
from app.services.datasource_mapping import (
    MappingError,
    build_heuristic_mapping,
    execute_mapping,
    redact_for_llm,
    stable_payload_hash,
)
from app.services.custom_datasource_runtime import (
    CustomDatasourceRuntimeError,
    fetch_rest_payload,
    get_custom_stream_status,
    run_mapped_rest_config,
    run_mapped_websocket_config,
    start_custom_stream,
    stop_custom_stream,
    test_websocket_config,
)

DATASOURCE_MAPPING_PROMPT_KEY = "datasource.mapping"
from app.services.datasource_connectivity import (
    _resolve_aisstream_api_key,
    _resolve_spacetrack_credentials_with_override,
    get_builtin_connection_status,
    save_connectivity_success,
    strip_connectivity_validation,
    test_builtin_connectivity,
)
from app.services.barentswatch import resolve_barentswatch_config
from app.services.persistent_logs import record_audit_log

router = APIRouter()

SECRET_REVEAL_ROLES = {"admin", "super_admin"}


def _user_role_value(user: User) -> str:
    role = getattr(user, "role", "")
    return str(getattr(role, "value", role) or "").lower()


def _user_display_name(user: User) -> str:
    return str(getattr(user, "username", None) or getattr(user, "email", None) or getattr(user, "id", ""))


async def _record_datasource_secret_reveal(
    *,
    current_user: User,
    request: Request,
    target_id: str,
    result: str,
    details: dict[str, Any],
) -> None:
    await record_audit_log(
        action="datasource_config.secret.reveal",
        actor_id=getattr(current_user, "id", None),
        actor_name=_user_display_name(current_user),
        target_type="datasource_config_secret",
        target_id=target_id,
        result=result,
        ip=request.client.host if request.client else None,
        details=details,
    )


async def _ensure_datasource_secret_reveal_allowed(
    current_user: User,
    request: Request,
    target_id: str,
    details: dict[str, Any],
) -> None:
    if _user_role_value(current_user) in SECRET_REVEAL_ROLES:
        return
    await _record_datasource_secret_reveal(
        current_user=current_user,
        request=request,
        target_id=target_id,
        result="denied",
        details={**details, "role": _user_role_value(current_user)},
    )
    raise HTTPException(
        status_code=status.HTTP_403_FORBIDDEN,
        detail="Only administrators can reveal datasource credentials",
    )

def _default_builtin_config(name: str) -> dict[str, Any]:
    return {"timeout": 30, "retry": 3}


def _default_builtin_source_type(name: str) -> str:
    if name == "aisstream_vessels":
        return "websocket"
    return "http"


class DataSourceConfigCreate(BaseModel):
    name: str = Field(..., min_length=1, max_length=100)
    description: Optional[str] = None
    source_type: str = Field(..., description="rest, websocket, http, api, database")
    endpoint: str = Field(..., max_length=500)
    auth_type: str = Field(default="none", description="none, bearer, api_key, basic")
    auth_config: dict = Field(default={})
    headers: dict = Field(default={})
    config: dict = Field(default={"timeout": 30, "retry": 3})


class DataSourceConfigUpdate(BaseModel):
    name: Optional[str] = Field(None, min_length=1, max_length=100)
    description: Optional[str] = None
    source_type: Optional[str] = None
    endpoint: Optional[str] = Field(None, max_length=500)
    auth_type: Optional[str] = None
    auth_config: Optional[dict] = None
    headers: Optional[dict] = None
    config: Optional[dict] = None
    is_active: Optional[bool] = None


class DataSourceConfigResponse(BaseModel):
    id: int
    name: str
    description: Optional[str]
    source_type: str
    endpoint: str
    auth_type: str
    headers: dict
    config: dict
    is_active: bool
    created_at: datetime
    updated_at: datetime

    class Config:
        from_attributes = True


def _is_builtin_config_name(name: str | None) -> bool:
    return bool(name and name in DEFAULT_DATASOURCES)


async def _ensure_builtin_connection_verified(
    db: AsyncSession,
    config_data: DataSourceConfigCreate,
) -> None:
    if not _is_builtin_config_name(config_data.name):
        return

    status_result = await get_builtin_connection_status(
        db,
        config_data.name,
        config_data.endpoint,
        config_data.auth_type,
        config_data.headers,
        config_data.config,
    )
    if not status_result.get("connected"):
        raise HTTPException(
            status_code=400,
            detail=status_result.get("message") or "请先完成连接验证，再保存内置采集器配置。",
        )


class CustomSampleRequest(BaseModel):
    datasource_config_id: Optional[int] = None
    config: Optional[DataSourceConfigCreate] = None
    limit_bytes: int = Field(default=200000, ge=1000, le=1000000)


class MappingProposeRequest(BaseModel):
    sample_payload: Any
    target_schema: str
    use_ai: bool = True


class MappingPreviewRequest(BaseModel):
    sample_payload: Any
    target_schema: str
    mapping_json: dict
    limit: int = Field(default=20, ge=1, le=100)


class MappingTemplateCreate(BaseModel):
    datasource_config_id: int
    target_schema: str
    mapping_json: dict
    sample_payload: Any | None = None
    sample_payload_hash: Optional[str] = None
    validation_status: str = Field(default="draft", pattern="^(draft|valid|invalid)$")
    is_active: bool = False


class MappingTemplateUpdate(BaseModel):
    target_schema: Optional[str] = None
    mapping_json: Optional[dict] = None
    sample_payload: Any | None = None
    sample_payload_hash: Optional[str] = None
    validation_status: Optional[str] = Field(default=None, pattern="^(draft|valid|invalid)$")
    is_active: Optional[bool] = None


async def test_endpoint(
    endpoint: str,
    auth_type: str,
    auth_config: dict,
    headers: dict,
    config: dict,
) -> dict:
    """Test an endpoint connection"""
    timeout = config.get("timeout", 30)
    test_headers = headers.copy()

    # Add auth headers
    if auth_type == "bearer" and auth_config.get("token"):
        test_headers["Authorization"] = f"Bearer {auth_config['token']}"
    elif auth_type == "api_key" and auth_config.get("api_key"):
        key_name = auth_config.get("key_name", "X-API-Key")
        test_headers[key_name] = auth_config["api_key"]
    elif auth_type == "basic":
        username = auth_config.get("username", "")
        password = auth_config.get("password", "")
        credentials = f"{username}:{password}"
        encoded = base64.b64encode(credentials.encode()).decode()
        test_headers["Authorization"] = f"Basic {encoded}"

    async with httpx.AsyncClient(timeout=timeout) as client:
        response = await client.get(endpoint, headers=test_headers)
        response.raise_for_status()
        return {
            "status_code": response.status_code,
            "success": True,
            "response_time_ms": response.elapsed.total_seconds() * 1000,
            "data_preview": str(response.json()[:3])
            if response.headers.get("content-type", "").startswith("application/json")
            else response.text[:200],
        }


def _build_request_headers(auth_type: str, auth_config: dict, headers: dict) -> dict[str, str]:
    request_headers = {str(key): str(value) for key, value in (headers or {}).items()}
    auth_type = str(auth_type or "none").lower()
    auth_config = auth_config or {}

    if auth_type == "bearer" and auth_config.get("token"):
        request_headers["Authorization"] = f"Bearer {auth_config['token']}"
    elif auth_type == "api_key" and auth_config.get("api_key"):
        location = str(auth_config.get("in") or auth_config.get("location") or "header").lower()
        if location != "query":
            key_name = auth_config.get("key_name", "X-API-Key")
            request_headers[str(key_name)] = str(auth_config["api_key"])
    elif auth_type == "basic":
        username = auth_config.get("username", "")
        password = auth_config.get("password", "")
        credentials = f"{username}:{password}"
        encoded = base64.b64encode(credentials.encode()).decode()
        request_headers["Authorization"] = f"Basic {encoded}"
    return request_headers


def _build_query_params(auth_type: str, auth_config: dict, config: dict) -> dict[str, Any]:
    params = {}
    candidate = (config or {}).get("params") or (config or {}).get("query_params")
    if isinstance(candidate, dict):
        params.update(candidate)

    auth_type = str(auth_type or "none").lower()
    auth_config = auth_config or {}
    if auth_type == "api_key" and auth_config.get("api_key"):
        location = str(auth_config.get("in") or auth_config.get("location") or "header").lower()
        if location == "query":
            key_name = auth_config.get("key_name") or auth_config.get("param_name") or "api_key"
            params[str(key_name)] = auth_config["api_key"]
    return params


async def fetch_custom_sample_from_config(config: DataSourceConfig, limit_bytes: int) -> Any:
    if str(config.source_type or "").lower() in {"websocket", "ws"}:
        raise HTTPException(status_code=400, detail="WebSocket sources must use connection test or run-mapped stream.")
    request_config = config.config or {}
    method = str(request_config.get("method") or request_config.get("request_method") or "GET").upper()
    if method not in {"GET", "POST"}:
        raise HTTPException(status_code=400, detail="Only GET and POST sample requests are supported.")

    headers = _build_request_headers(config.auth_type, config.auth_config or {}, config.headers or {})
    params = _build_query_params(config.auth_type, config.auth_config or {}, request_config)
    timeout = float(request_config.get("timeout", 30))
    json_body = request_config.get("json_body")
    if json_body is None and str(request_config.get("body_type") or "").lower() in {"json", ""}:
        candidate = request_config.get("body")
        if isinstance(candidate, (dict, list)):
            json_body = candidate

    async with httpx.AsyncClient(timeout=timeout, follow_redirects=True) as client:
        response = await client.request(
            method,
            config.endpoint,
            headers=headers,
            params=params or None,
            json=json_body,
        )
        response.raise_for_status()
        content = response.content[:limit_bytes]
        if "application/json" in response.headers.get("content-type", ""):
            return json.loads(content.decode(response.encoding or "utf-8"))
        return {"text": content.decode(response.encoding or "utf-8", errors="replace")}


def _parse_mapping_from_ai_text(content: str) -> dict[str, Any] | None:
    if not content:
        return None

    candidates = [content]
    fenced = re.findall(r"```(?:json)?\s*(\{.*?\})\s*```", content, flags=re.DOTALL)
    candidates = fenced + candidates
    for candidate in candidates:
        try:
            parsed = json.loads(candidate)
        except json.JSONDecodeError:
            continue
        if isinstance(parsed, dict) and isinstance(parsed.get("fields"), dict):
            return parsed
    return None


async def _get_config_for_sample(
    payload: CustomSampleRequest,
    db: AsyncSession,
) -> DataSourceConfig:
    if payload.datasource_config_id is not None:
        result = await db.execute(
            select(DataSourceConfig).where(DataSourceConfig.id == payload.datasource_config_id)
        )
        config = result.scalar_one_or_none()
        if not config:
            raise HTTPException(status_code=404, detail="Configuration not found")
        return config

    if payload.config is None:
        raise HTTPException(status_code=400, detail="datasource_config_id or config is required")

    config_data = payload.config
    return DataSourceConfig(
        name=config_data.name,
        description=config_data.description,
        source_type=config_data.source_type,
        endpoint=config_data.endpoint,
        auth_type=config_data.auth_type,
        auth_config=config_data.auth_config,
        headers=config_data.headers,
        config=config_data.config,
    )


def serialize_mapping_template(template: DataSourceMappingTemplate) -> dict[str, Any]:
    return {
        "id": template.id,
        "datasource_config_id": template.datasource_config_id,
        "target_schema": template.target_schema,
        "mapping_json": template.mapping_json,
        "sample_payload_hash": template.sample_payload_hash,
        "validation_status": template.validation_status,
        "version": template.version,
        "is_active": template.is_active,
        "created_at": to_iso8601_utc(template.created_at),
        "updated_at": to_iso8601_utc(template.updated_at),
    }


@router.get("/configs")
async def list_configs(
    active_only: bool = False,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    """List all user-defined data source configurations"""
    query = select(DataSourceConfig)
    if active_only:
        query = query.where(DataSourceConfig.is_active)
    query = query.order_by(DataSourceConfig.created_at.desc())

    result = await db.execute(query)
    configs = result.scalars().all()

    return {
        "total": len(configs),
        "data": [
            {
                "id": c.id,
                "name": c.name,
                "description": c.description,
                "source_type": c.source_type,
                "endpoint": c.endpoint,
                "auth_type": c.auth_type,
                "headers": c.headers,
                "config": c.config,
                "is_active": c.is_active,
                "created_at": to_iso8601_utc(c.created_at),
                "updated_at": to_iso8601_utc(c.updated_at),
            }
            for c in configs
        ],
    }


@router.get("/configs/all")
async def list_all_datasources(
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    """List all data sources: YAML defaults + DB overrides"""
    from app.core.data_sources import get_data_sources_config

    config = get_data_sources_config()

    db_query = await db.execute(select(DataSourceConfig))
    db_configs = {c.name: c for c in db_query.scalars().all()}

    result = []
    for name, metadata in DEFAULT_DATASOURCES.items():
        yaml_url = config.get_yaml_url(name)
        db_config = db_configs.get(name)
        default_config = _default_builtin_config(name)
        default_url = yaml_url
        db_auth_config = db_config.auth_config or {} if db_config else {}

        result.append(
            {
                "name": name,
                "requires_credentials": bool(metadata.get("requires_credentials", False)),
                "credential_provider": metadata.get("credential_provider"),
                "credential_status": metadata.get("credential_status", "none"),
                "default_url": default_url,
                "endpoint": db_config.endpoint if db_config else default_url,
                "is_overridden": db_config is not None and db_config.endpoint != yaml_url
                if default_url
                else db_config is not None,
                "is_active": db_config.is_active if db_config else True,
                "source_type": db_config.source_type if db_config else _default_builtin_source_type(name),
                "auth_type": db_config.auth_type if db_config else "none",
                "auth_config": {
                    "client_id": db_auth_config.get("client_id") or "",
                    "username": db_auth_config.get("username") or "",
                    "key_name": db_auth_config.get("key_name") or db_auth_config.get("param_name") or "",
                    "param_name": db_auth_config.get("param_name") or db_auth_config.get("key_name") or "",
                    "location": db_auth_config.get("location") or db_auth_config.get("in") or "",
                    "in": db_auth_config.get("in") or db_auth_config.get("location") or "",
                },
                "auth_configured": {
                    "api_key": bool(db_auth_config.get("api_key")),
                    "client_id": bool(db_auth_config.get("client_id")),
                    "client_secret": bool(db_auth_config.get("client_secret")),
                    "username": bool(db_auth_config.get("username")),
                    "password": bool(db_auth_config.get("password")),
                },
                "headers": db_config.headers if db_config else {},
                "config": strip_connectivity_validation(db_config.config if db_config else default_config),
                "config_id": db_config.id if db_config else None,
                "description": db_config.description
                if db_config
                else f"内置采集器默认配置：{metadata.get('display_name') or metadata.get('name') or name}",
            }
        )

    return {"total": len(result), "data": result}


@router.get("/configs/secrets")
async def reveal_builtin_config_secrets(
    request: Request,
    name: str = Query(..., min_length=1),
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    """Reveal configured built-in datasource credentials for admin editing."""
    source = name.strip()
    metadata = DEFAULT_DATASOURCES.get(source)
    if not metadata or not metadata.get("requires_credentials"):
        raise HTTPException(status_code=404, detail="Credentialed datasource config not found")

    provider = str(metadata.get("credential_provider") or "")
    target_id = f"datasource_config:{source}"
    await _ensure_datasource_secret_reveal_allowed(
        current_user,
        request,
        target_id,
        {"source": source, "provider": provider},
    )

    result = await db.execute(select(DataSourceConfig).where(DataSourceConfig.name == source))
    record = result.scalar_one_or_none()
    auth_config = dict(record.auth_config or {}) if record else {}
    payload: dict[str, Any] = {
        "name": source,
        "provider": provider,
    }
    details: dict[str, Any] = {"source": source, "provider": provider}

    if provider == "barentswatch":
        resolved = await resolve_barentswatch_config(db)
        client_id = str(auth_config.get("client_id") or resolved.client_id or "")
        client_secret = str(auth_config.get("client_secret") or resolved.client_secret or "")
        source_label = "datasource_config" if auth_config.get("client_id") or auth_config.get("client_secret") else resolved.credential_source
        payload.update(
            {
                "client_id": client_id,
                "client_secret": client_secret,
                "client_id_source": source_label if client_id else "missing",
                "client_secret_source": source_label if client_secret else "missing",
            }
        )
        details.update(
            {
                "client_id_configured": bool(client_id),
                "client_secret_configured": bool(client_secret),
                "credential_source": source_label,
            }
        )
    elif provider == "aisstream":
        api_key, api_key_source = await _resolve_aisstream_api_key(db)
        payload.update({"api_key": api_key, "api_key_source": api_key_source})
        details.update({"api_key_configured": bool(api_key), "api_key_source": api_key_source})
    elif provider == "spacetrack":
        if auth_config.get("username") or auth_config.get("password"):
            username = str(auth_config.get("username") or "")
            password = str(auth_config.get("password") or "")
            credential_source = "datasource_config"
        else:
            username, password, credential_source = _resolve_spacetrack_credentials_with_override()
        payload.update(
            {
                "username": username,
                "password": password,
                "username_source": credential_source if username else "missing",
                "password_source": credential_source if password else "missing",
            }
        )
        details.update(
            {
                "username_configured": bool(username),
                "password_configured": bool(password),
                "credential_source": credential_source,
            }
        )
    else:
        raise HTTPException(status_code=400, detail="Datasource credential provider is not supported")

    await _record_datasource_secret_reveal(
        current_user=current_user,
        request=request,
        target_id=target_id,
        result="success",
        details=details,
    )
    return payload


@router.get("/configs/{config_id}")
async def get_config(
    config_id: int,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    """Get a single data source configuration"""
    result = await db.execute(select(DataSourceConfig).where(DataSourceConfig.id == config_id))
    config = result.scalar_one_or_none()

    if not config:
        raise HTTPException(status_code=404, detail="Configuration not found")

    return {
        "id": config.id,
        "name": config.name,
        "description": config.description,
        "source_type": config.source_type,
        "endpoint": config.endpoint,
        "auth_type": config.auth_type,
        "auth_config": {},  # Don't return sensitive data
        "headers": config.headers,
        "config": config.config,
        "is_active": config.is_active,
        "created_at": to_iso8601_utc(config.created_at),
        "updated_at": to_iso8601_utc(config.updated_at),
    }


@router.post("/configs")
async def create_config(
    config_data: DataSourceConfigCreate,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    """Create a new data source configuration"""
    config = DataSourceConfig(
        name=config_data.name,
        description=config_data.description,
        source_type=config_data.source_type,
        endpoint=config_data.endpoint,
        auth_type=config_data.auth_type,
        auth_config=config_data.auth_config,
        headers=config_data.headers,
        config=strip_connectivity_validation(config_data.config),
    )

    db.add(config)
    await db.commit()
    await db.refresh(config)

    cache.delete_pattern("datasource_configs:*")

    return {
        "id": config.id,
        "name": config.name,
        "message": "Configuration created successfully",
    }


@router.put("/configs/{config_id}")
async def update_config(
    config_id: int,
    config_data: DataSourceConfigUpdate,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    """Update a data source configuration"""
    result = await db.execute(select(DataSourceConfig).where(DataSourceConfig.id == config_id))
    config = result.scalar_one_or_none()

    if not config:
        raise HTTPException(status_code=404, detail="Configuration not found")

    update_data = config_data.model_dump(exclude_unset=True)
    for field, value in update_data.items():
        if field == "config":
            value = strip_connectivity_validation(value)
        if field == "auth_config" and value == {} and (config.auth_config or {}):
            continue
        setattr(config, field, value)

    await db.commit()
    await db.refresh(config)

    cache.delete_pattern("datasource_configs:*")

    return {
        "id": config.id,
        "name": config.name,
        "message": "Configuration updated successfully",
    }


@router.delete("/configs/{config_id}")
async def delete_config(
    config_id: int,
    delete_mappings: bool = Query(False),
    delete_source_data: bool = Query(False),
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    """Delete a data source configuration"""
    result = await db.execute(select(DataSourceConfig).where(DataSourceConfig.id == config_id))
    config = result.scalar_one_or_none()

    if not config:
        raise HTTPException(status_code=404, detail="Configuration not found")

    deleted_mappings = 0
    deleted_records = {
        "collected_data": 0,
        "ais_raw_observations": 0,
        "ais_source_health": 0,
    }

    if delete_source_data:
        collected_result = await db.execute(
            delete(CollectedData).where(CollectedData.source == config.name)
        )
        raw_result = await db.execute(
            delete(AISRawObservation).where(AISRawObservation.source == config.name)
        )
        health_result = await db.execute(
            delete(AISSourceHealth).where(AISSourceHealth.source == config.name)
        )
        deleted_records = {
            "collected_data": collected_result.rowcount or 0,
            "ais_raw_observations": raw_result.rowcount or 0,
            "ais_source_health": health_result.rowcount or 0,
        }

    if delete_mappings or delete_source_data:
        mapping_result = await db.execute(
            delete(DataSourceMappingTemplate).where(
                DataSourceMappingTemplate.datasource_config_id == config_id
            )
        )
        deleted_mappings = mapping_result.rowcount or 0

    await db.delete(config)
    await db.commit()

    cache.delete_pattern("datasource_configs:*")

    if delete_source_data and (config.config or {}).get("target_schema") == "vessel_ais":
        from app.core.websocket.broadcaster import broadcaster

        await broadcaster.broadcast_custom(
            "vessels",
            {
                "action": "reload",
                "source": config.name,
                "reason": "custom_source_deleted",
            },
        )

    return {
        "message": "Configuration deleted successfully",
        "deleted_mappings": deleted_mappings,
        "deleted_records": deleted_records,
    }


@router.post("/configs/{config_id}/test")
async def test_config(
    config_id: int,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    """Test a data source configuration"""
    result = await db.execute(select(DataSourceConfig).where(DataSourceConfig.id == config_id))
    config = result.scalar_one_or_none()

    if not config:
        raise HTTPException(status_code=404, detail="Configuration not found")

    try:
        if str(config.source_type or "").lower() in {"websocket", "ws"}:
            return await test_websocket_config(config)
        result = await test_endpoint(
            endpoint=config.endpoint,
            auth_type=config.auth_type,
            auth_config=config.auth_config or {},
            headers=config.headers or {},
            config=config.config or {},
        )
        return result

    except httpx.HTTPStatusError as e:
        return {
            "success": False,
            "error": f"HTTP Error: {e.response.status_code}",
            "message": str(e),
        }
    except Exception as e:
        return {
            "success": False,
            "error": "Connection failed",
            "message": str(e),
        }


@router.post("/configs/test")
async def test_new_config(
    config_data: DataSourceConfigCreate,
    current_user: User = Depends(get_current_user),
):
    """Test a new data source configuration without saving"""
    try:
        if str(config_data.source_type or "").lower() in {"websocket", "ws"}:
            config = DataSourceConfig(
                name=config_data.name,
                description=config_data.description,
                source_type=config_data.source_type,
                endpoint=config_data.endpoint,
                auth_type=config_data.auth_type,
                auth_config=config_data.auth_config,
                headers=config_data.headers,
                config=config_data.config,
            )
            return await test_websocket_config(config)
        result = await test_endpoint(
            endpoint=config_data.endpoint,
            auth_type=config_data.auth_type,
            auth_config=config_data.auth_config or {},
            headers=config_data.headers or {},
            config=config_data.config or {},
        )
        return result

    except httpx.HTTPStatusError as e:
        return {
            "success": False,
            "error": f"HTTP Error: {e.response.status_code}",
            "message": str(e),
        }
    except Exception as e:
        return {
            "success": False,
            "error": "Connection failed",
            "message": str(e),
        }


@router.post("/configs/builtin/connection-status")
async def get_builtin_config_connection_status(
    config_data: DataSourceConfigCreate,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    if not _is_builtin_config_name(config_data.name):
        raise HTTPException(status_code=400, detail="Only built-in datasource configs are supported.")

    return await get_builtin_connection_status(
        db,
        config_data.name,
        config_data.endpoint,
        config_data.auth_type,
        config_data.headers,
        config_data.config,
    )


@router.post("/configs/builtin/connect")
async def connect_builtin_config(
    config_data: DataSourceConfigCreate,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    if not _is_builtin_config_name(config_data.name):
        raise HTTPException(status_code=400, detail="Only built-in datasource configs are supported.")

    result = await test_builtin_connectivity(
        config_data.name,
        config_data.endpoint,
        config_data.auth_type,
        config_data.headers,
        config_data.config,
        db,
        config_data.auth_config,
    )
    if result.get("success") and result.get("checksum"):
        validation = await save_connectivity_success(
            db,
            config_data.name,
            result["checksum"],
            result,
            connected_by="connection_button",
        )
        await db.commit()
        return {
            **result,
            "connected": True,
            "validation": validation,
        }

    return {
        **result,
        "connected": False,
    }


@router.post("/custom/sample")
async def fetch_custom_sample(
    payload: CustomSampleRequest,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    """Fetch a sample payload for a saved or draft custom data source."""
    config = await _get_config_for_sample(payload, db)
    try:
        sample = await fetch_custom_sample_from_config(config, payload.limit_bytes)
    except httpx.HTTPStatusError as exc:
        raise HTTPException(
            status_code=exc.response.status_code,
            detail=f"Sample request failed: HTTP {exc.response.status_code}",
        ) from exc
    except httpx.HTTPError as exc:
        raise HTTPException(status_code=502, detail=f"Sample request failed: {exc}") from exc

    return {
        "success": True,
        "sample_payload": sample,
        "sample_payload_hash": stable_payload_hash(sample),
        "redacted_preview": redact_for_llm(sample),
    }


@router.get("/target-schemas")
async def get_datasource_target_schemas(
    current_user: User = Depends(get_current_user),
):
    """List target schemas available for custom datasource mapping."""
    return {"data": list_target_schemas()}


@router.post("/mappings/propose")
async def propose_datasource_mapping(
    payload: MappingProposeRequest,
    current_user: User = Depends(get_current_user),
    ai_client: AIProviderClient = Depends(get_ai_provider_client),
):
    """Generate a mapping draft for a sample payload and target schema."""
    schema = get_target_schema(payload.target_schema)
    redacted_sample = redact_for_llm(payload.sample_payload)
    fallback_mapping = build_heuristic_mapping(redacted_sample, payload.target_schema)

    ai_error: str | None = None
    mapping = fallback_mapping
    generated_by = "heuristic"
    if payload.use_ai:
        try:
            prompt = await get_effective_prompt(db, DATASOURCE_MAPPING_PROMPT_KEY)
            response = await ai_client.analyze(
                SituationalAnalysisRequest(
                    title=f"Generate datasource mapping for {schema.key}",
                    objective=prompt.prompt,
                    system_prompt=prompt.system_prompt or None,
                    context={
                        "target_schema": schema.to_dict(),
                        "sample_payload": redacted_sample,
                        "mapping_dsl_example": fallback_mapping,
                    },
                    observations=[
                        "Use JSONPath-like paths beginning with $.",
                        "Never generate executable code.",
                        "Use field types from the target schema.",
                    ],
                    constraints=[
                        "Return a single JSON object.",
                        "Do not include credentials or secrets.",
                        "Mark uncertain optional fields with default null.",
                    ],
                )
            )
            parsed = _parse_mapping_from_ai_text(response.content)
            if parsed:
                mapping = parsed
                generated_by = "ai_provider"
            else:
                ai_error = "AI provider did not return a valid mapping JSON object."
        except HTTPException as exc:
            ai_error = str(exc.detail)

    mapping.setdefault("meta", {})
    if isinstance(mapping["meta"], dict):
        mapping["meta"].update(
            {
                "generated_by": generated_by,
                "requires_review": True,
                "ai_error": ai_error,
            }
        )

    return {
        "target_schema": schema.to_dict(),
        "mapping_json": mapping,
        "sample_payload_hash": stable_payload_hash(payload.sample_payload),
        "redacted_sample_payload": redacted_sample,
    }


@router.post("/mappings/preview")
async def preview_datasource_mapping(
    payload: MappingPreviewRequest,
    current_user: User = Depends(get_current_user),
):
    """Preview deterministic mapping output for a sample payload."""
    try:
        preview = execute_mapping(
            payload.sample_payload,
            payload.mapping_json,
            payload.target_schema,
            limit=payload.limit,
        )
    except (MappingError, ValueError) as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc

    return {
        "success": preview["failed_count"] == 0,
        "preview": preview,
        "sample_payload_hash": stable_payload_hash(payload.sample_payload),
    }


@router.get("/mappings")
async def list_datasource_mappings(
    datasource_config_id: Optional[int] = None,
    active_only: bool = False,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    """List saved mapping templates."""
    query = select(DataSourceMappingTemplate).order_by(
        DataSourceMappingTemplate.datasource_config_id,
        DataSourceMappingTemplate.version.desc(),
    )
    if datasource_config_id is not None:
        query = query.where(DataSourceMappingTemplate.datasource_config_id == datasource_config_id)
    if active_only:
        query = query.where(DataSourceMappingTemplate.is_active.is_(True))

    result = await db.execute(query)
    mappings = result.scalars().all()
    return {"total": len(mappings), "data": [serialize_mapping_template(item) for item in mappings]}


@router.post("/mappings")
async def create_datasource_mapping(
    payload: MappingTemplateCreate,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    """Save a mapping template for a datasource config."""
    get_target_schema(payload.target_schema)
    datasource = await db.get(DataSourceConfig, payload.datasource_config_id)
    if not datasource:
        raise HTTPException(status_code=404, detail="Configuration not found")

    if payload.sample_payload is not None:
        try:
            execute_mapping(payload.sample_payload, payload.mapping_json, payload.target_schema, limit=100)
        except (MappingError, ValueError) as exc:
            raise HTTPException(status_code=400, detail=f"Mapping validation failed: {exc}") from exc

    result = await db.execute(
        select(func.max(DataSourceMappingTemplate.version)).where(
            DataSourceMappingTemplate.datasource_config_id == payload.datasource_config_id,
            DataSourceMappingTemplate.target_schema == payload.target_schema,
        )
    )
    next_version = int(result.scalar() or 0) + 1

    if payload.is_active:
        await db.execute(
            DataSourceMappingTemplate.__table__.update()
            .where(DataSourceMappingTemplate.datasource_config_id == payload.datasource_config_id)
            .values(is_active=False)
        )

    template = DataSourceMappingTemplate(
        datasource_config_id=payload.datasource_config_id,
        target_schema=payload.target_schema,
        mapping_json=payload.mapping_json,
        sample_payload_hash=payload.sample_payload_hash
        or (stable_payload_hash(payload.sample_payload) if payload.sample_payload is not None else None),
        validation_status=payload.validation_status,
        version=next_version,
        is_active=payload.is_active,
    )
    db.add(template)
    await db.commit()
    await db.refresh(template)
    return {"message": "Mapping template saved successfully", "data": serialize_mapping_template(template)}


@router.put("/mappings/{mapping_id}")
async def update_datasource_mapping(
    mapping_id: int,
    payload: MappingTemplateUpdate,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    """Update a mapping template in place."""
    template = await db.get(DataSourceMappingTemplate, mapping_id)
    if not template:
        raise HTTPException(status_code=404, detail="Mapping template not found")

    target_schema = payload.target_schema or template.target_schema
    mapping_json = payload.mapping_json or template.mapping_json
    get_target_schema(target_schema)
    if payload.sample_payload is not None:
        try:
            execute_mapping(payload.sample_payload, mapping_json, target_schema, limit=100)
        except (MappingError, ValueError) as exc:
            raise HTTPException(status_code=400, detail=f"Mapping validation failed: {exc}") from exc

    if payload.is_active is True:
        await db.execute(
            DataSourceMappingTemplate.__table__.update()
            .where(DataSourceMappingTemplate.datasource_config_id == template.datasource_config_id)
            .where(DataSourceMappingTemplate.id != template.id)
            .values(is_active=False)
        )

    template.target_schema = target_schema
    template.mapping_json = mapping_json
    if payload.sample_payload_hash is not None:
        template.sample_payload_hash = payload.sample_payload_hash
    elif payload.sample_payload is not None:
        template.sample_payload_hash = stable_payload_hash(payload.sample_payload)
    if payload.validation_status is not None:
        template.validation_status = payload.validation_status
    if payload.is_active is not None:
        template.is_active = payload.is_active

    await db.commit()
    await db.refresh(template)
    return {"message": "Mapping template updated successfully", "data": serialize_mapping_template(template)}


@router.post("/{config_id}/run-mapped")
async def run_mapped_datasource(
    config_id: int,
    background: bool = Query(False, description="For WebSocket sources, start a background stream task."),
    debug_max_messages: int | None = Query(None, ge=1),
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    """Run a saved custom datasource through its active deterministic mapping."""
    datasource = await db.get(DataSourceConfig, config_id)
    if not datasource:
        raise HTTPException(status_code=404, detail="Configuration not found")

    try:
        if str(datasource.source_type or "").lower() in {"websocket", "ws"}:
            if background and debug_max_messages is None:
                started = start_custom_stream(config_id)
                if not started:
                    raise HTTPException(status_code=409, detail="Custom WebSocket source is already running")
                return {
                    "status": "started",
                    "datasource_config_id": config_id,
                    "stream": get_custom_stream_status(config_id),
                }
            return await run_mapped_websocket_config(
                db,
                datasource,
                debug_max_messages=debug_max_messages,
            )

        return await run_mapped_rest_config(db, datasource)
    except httpx.HTTPStatusError as exc:
        raise HTTPException(
            status_code=exc.response.status_code,
            detail=f"Datasource request failed: HTTP {exc.response.status_code}",
        ) from exc
    except httpx.HTTPError as exc:
        raise HTTPException(status_code=502, detail=f"Datasource request failed: {exc}") from exc
    except (CustomDatasourceRuntimeError, MappingError, ValueError) as exc:
        raise HTTPException(status_code=400, detail=f"Mapping failed: {exc}") from exc


@router.post("/{config_id}/stop-mapped")
async def stop_mapped_datasource(
    config_id: int,
    current_user: User = Depends(get_current_user),
):
    stopped = await stop_custom_stream(config_id)
    return {
        "status": "stopped" if stopped else "not_running",
        "datasource_config_id": config_id,
        "stream": get_custom_stream_status(config_id),
    }


@router.get("/{config_id}/stream-status")
async def get_mapped_stream_status(
    config_id: int,
    current_user: User = Depends(get_current_user),
):
    return get_custom_stream_status(config_id)