"""Runtime helpers for mapped custom data sources."""

from __future__ import annotations

import asyncio
import base64
import json
from datetime import UTC, datetime
from typing import Any

import httpx
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from app.core.target_schema_registry import TARGET_SCHEMAS
from app.db.session import async_session_factory
from app.models.datasource_config import DataSourceConfig
from app.models.datasource_mapping import DataSourceMappingTemplate
from app.services.datasource_mapping import (
    MappingError,
    execute_mapping,
    extract_path,
    persist_mapped_records,
)

DEFAULT_MAPPING_TEMPLATES: dict[str, dict[str, Any]] = {
    "vessel_ais": {
        "source": {"items_path": "$"},
        "fields": {
            "mmsi": {"path": "$.mmsi", "type": "integer"},
            "name": {"path": "$.name", "type": "string", "default": None},
            "lat": {"path": "$.lat", "type": "float"},
            "lon": {"path": "$.lon", "type": "float"},
            "sog": {"path": "$.sog", "type": "float", "default": None},
            "cog": {"path": "$.cog", "type": "float", "default": None},
            "heading": {"path": "$.heading", "type": "integer", "default": None},
            "nav_status": {"path": "$.nav_status", "type": "integer", "default": None},
            "callsign": {"path": "$.callsign", "type": "string", "default": None},
            "vessel_type": {"path": "$.vessel_type", "type": "string", "default": None},
            "vessel_type_name": {"path": "$.vessel_type_name", "type": "string", "default": None},
            "received_at": {"path": "$.received_at", "type": "datetime", "default": None},
        },
        "meta": {"generated_by": "default_template", "requires_review": False},
    },
}

RUNNING_CUSTOM_STREAM_TASKS: dict[int, asyncio.Task[Any]] = {}


class CustomDatasourceRuntimeError(RuntimeError):
    """Raised when a custom datasource cannot run."""


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: dict[str, Any] = {}
    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 load_active_mapping(
    db: AsyncSession,
    datasource_config_id: int,
) -> DataSourceMappingTemplate:
    result = await db.execute(
        select(DataSourceMappingTemplate)
        .where(DataSourceMappingTemplate.datasource_config_id == datasource_config_id)
        .where(DataSourceMappingTemplate.is_active.is_(True))
        .order_by(DataSourceMappingTemplate.version.desc())
        .limit(1)
    )
    mapping = result.scalar_one_or_none()
    if mapping is not None:
        return mapping

    datasource = await db.get(DataSourceConfig, datasource_config_id)
    if datasource is None:
        raise CustomDatasourceRuntimeError("Configuration not found")
    target_schema = (datasource.config or {}).get("target_schema")
    template_body = DEFAULT_MAPPING_TEMPLATES.get(str(target_schema or "")) if target_schema else None
    if not template_body or target_schema not in TARGET_SCHEMAS:
        raise CustomDatasourceRuntimeError(
            "No active mapping template found and no default template available for this target schema"
        )

    mapping = DataSourceMappingTemplate(
        datasource_config_id=datasource_config_id,
        target_schema=str(target_schema),
        mapping_json=template_body,
        sample_payload_hash=None,
        validation_status="valid",
        version=1,
        is_active=True,
    )
    db.add(mapping)
    await db.commit()
    await db.refresh(mapping)
    return mapping


async def fetch_rest_payload(config: DataSourceConfig, limit_bytes: int) -> Any:
    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 CustomDatasourceRuntimeError("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")}


async def run_mapped_rest_config(
    db: AsyncSession,
    datasource: DataSourceConfig,
) -> dict[str, Any]:
    mapping = await load_active_mapping(db, datasource.id)
    sample = await fetch_rest_payload(datasource, 5_000_000)
    mapped = execute_mapping(sample, mapping.mapping_json, mapping.target_schema)
    if mapped["failed_count"] > 0:
        return {
            "status": "failed",
            "datasource_config_id": datasource.id,
            "mapping_id": mapping.id,
            "mapping_version": mapping.version,
            "target_schema": mapping.target_schema,
            "mapped_count": mapped["mapped_count"],
            "failed_count": mapped["failed_count"],
            "errors": mapped["errors"][:20],
        }

    request_config = datasource.config or {}
    written_count = await persist_mapped_records(
        db,
        datasource_name=datasource.name,
        datasource_config_id=datasource.id,
        target_schema=mapping.target_schema,
        records=mapped["records"],
        mapping_version=mapping.version,
        delivery_mode=request_config.get("delivery_mode") or "polling",
        transport="http",
    )
    return {
        "status": "success",
        "datasource_config_id": datasource.id,
        "mapping_id": mapping.id,
        "mapping_version": mapping.version,
        "target_schema": mapping.target_schema,
        "fetched_count": mapped["total_items"],
        "mapped_count": mapped["mapped_count"],
        "written_count": written_count,
    }


def _items_from_ws_message(payload: Any, config: dict) -> Any:
    message_path = config.get("ws_message_path")
    items_path = config.get("ws_items_path")
    value = extract_path(payload, message_path) if message_path else payload
    return extract_path(value, items_path) if items_path else value


async def _connect_websocket(endpoint: str, headers: dict[str, str]):
    import websockets

    try:
        return await websockets.connect(endpoint, additional_headers=headers or None)
    except TypeError:
        return await websockets.connect(endpoint, extra_headers=headers or None)


async def test_websocket_config(config: DataSourceConfig) -> dict[str, Any]:
    if not str(config.endpoint or "").startswith(("ws://", "wss://")):
        raise CustomDatasourceRuntimeError("WebSocket datasource endpoint must start with ws:// or wss://")

    runtime_config = config.config or {}
    headers = build_request_headers(config.auth_type, config.auth_config or {}, config.headers or {})
    receive_timeout = float(runtime_config.get("receive_timeout_seconds") or runtime_config.get("timeout") or 10)
    async with await _connect_websocket(config.endpoint, headers) as websocket:
        subscribe_message = runtime_config.get("ws_subscribe_message")
        if isinstance(subscribe_message, (dict, list)):
            await websocket.send(json.dumps(subscribe_message))
        elif isinstance(subscribe_message, str) and subscribe_message.strip():
            await websocket.send(subscribe_message)
        raw_message = await asyncio.wait_for(websocket.recv(), timeout=receive_timeout)
        return {
            "success": True,
            "message_preview": raw_message[:1000] if isinstance(raw_message, str) else str(raw_message)[:1000],
        }


async def run_mapped_websocket_config(
    db: AsyncSession,
    datasource: DataSourceConfig,
    *,
    debug_max_messages: int | None = None,
    use_config_debug_max_messages: bool = True,
) -> dict[str, Any]:
    if not str(datasource.endpoint or "").startswith(("ws://", "wss://")):
        raise CustomDatasourceRuntimeError("WebSocket datasource endpoint must start with ws:// or wss://")

    mapping = await load_active_mapping(db, datasource.id)
    runtime_config = datasource.config or {}
    max_messages = debug_max_messages
    if max_messages is None and use_config_debug_max_messages:
        max_messages = runtime_config.get("debug_max_messages")
    max_messages = int(max_messages) if max_messages else None
    receive_timeout = float(runtime_config.get("receive_timeout_seconds") or runtime_config.get("timeout") or 30)
    reconnect = bool(runtime_config.get("ws_reconnect", True))
    reconnect_delay = float(runtime_config.get("reconnect_delay_seconds") or 3)
    headers = build_request_headers(datasource.auth_type, datasource.auth_config or {}, datasource.headers or {})

    messages_seen = 0
    mapped_count = 0
    failed_count = 0
    written_count = 0
    errors: list[dict[str, Any]] = []
    started_at = datetime.now(UTC)

    while True:
        try:
            async with await _connect_websocket(datasource.endpoint, headers) as websocket:
                subscribe_message = runtime_config.get("ws_subscribe_message")
                if isinstance(subscribe_message, (dict, list)):
                    await websocket.send(json.dumps(subscribe_message))
                elif isinstance(subscribe_message, str) and subscribe_message.strip():
                    await websocket.send(subscribe_message)

                while True:
                    raw_message = await asyncio.wait_for(websocket.recv(), timeout=receive_timeout)
                    messages_seen += 1
                    try:
                        payload = json.loads(raw_message)
                    except json.JSONDecodeError as exc:
                        failed_count += 1
                        errors.append({"message": "invalid_json", "error": str(exc)})
                        continue

                    extracted = _items_from_ws_message(payload, runtime_config)
                    try:
                        mapped = execute_mapping(extracted, mapping.mapping_json, mapping.target_schema)
                    except (MappingError, ValueError) as exc:
                        failed_count += 1
                        errors.append({"message": "mapping_failed", "error": str(exc)})
                        continue

                    mapped_count += mapped["mapped_count"]
                    failed_count += mapped["failed_count"]
                    if mapped["errors"]:
                        errors.extend(mapped["errors"][:5])
                    if mapped["records"]:
                        written_count += await persist_mapped_records(
                            db,
                            datasource_name=datasource.name,
                            datasource_config_id=datasource.id,
                            target_schema=mapping.target_schema,
                            records=mapped["records"],
                            mapping_version=mapping.version,
                            delivery_mode=runtime_config.get("delivery_mode") or "realtime_stream",
                            transport="websocket",
                        )

                    if max_messages and messages_seen >= max_messages:
                        return {
                            "status": "success",
                            "datasource_config_id": datasource.id,
                            "mapping_id": mapping.id,
                            "mapping_version": mapping.version,
                            "target_schema": mapping.target_schema,
                            "messages_seen": messages_seen,
                            "mapped_count": mapped_count,
                            "failed_count": failed_count,
                            "written_count": written_count,
                            "errors": errors[:20],
                            "execution_time_seconds": (datetime.now(UTC) - started_at).total_seconds(),
                        }
        except asyncio.CancelledError:
            raise
        except Exception as exc:
            failed_count += 1
            errors.append({"message": "websocket_error", "error": f"{exc.__class__.__name__}: {exc}"})
            if not reconnect or max_messages:
                return {
                    "status": "failed" if written_count == 0 else "partial",
                    "datasource_config_id": datasource.id,
                    "mapping_id": mapping.id,
                    "mapping_version": mapping.version,
                    "target_schema": mapping.target_schema,
                    "messages_seen": messages_seen,
                    "mapped_count": mapped_count,
                    "failed_count": failed_count,
                    "written_count": written_count,
                    "errors": errors[:20],
                }
            await asyncio.sleep(reconnect_delay)


async def run_custom_stream_by_id(config_id: int) -> dict[str, Any]:
    async with async_session_factory() as db:
        datasource = await db.get(DataSourceConfig, config_id)
        if not datasource:
            raise CustomDatasourceRuntimeError("Configuration not found")
        return await run_mapped_websocket_config(
            db,
            datasource,
            use_config_debug_max_messages=False,
        )


def start_custom_stream(config_id: int) -> bool:
    existing = RUNNING_CUSTOM_STREAM_TASKS.get(config_id)
    if existing is not None and not existing.done():
        return False
    task = asyncio.create_task(run_custom_stream_by_id(config_id), name=f"custom-stream:{config_id}")
    RUNNING_CUSTOM_STREAM_TASKS[config_id] = task

    def _cleanup(done_task: asyncio.Task[Any]) -> None:
        if RUNNING_CUSTOM_STREAM_TASKS.get(config_id) is done_task:
            RUNNING_CUSTOM_STREAM_TASKS.pop(config_id, None)

    task.add_done_callback(_cleanup)
    return True


async def stop_custom_stream(config_id: int) -> bool:
    task = RUNNING_CUSTOM_STREAM_TASKS.get(config_id)
    if task is None or task.done():
        RUNNING_CUSTOM_STREAM_TASKS.pop(config_id, None)
        return False
    task.cancel()
    try:
        await task
    except asyncio.CancelledError:
        return True
    return task.cancelled()


def get_custom_stream_status(config_id: int) -> dict[str, Any]:
    task = RUNNING_CUSTOM_STREAM_TASKS.get(config_id)
    return {
        "config_id": config_id,
        "running": bool(task and not task.done()),
        "done": bool(task and task.done()),
    }