from __future__ import annotations

import re
from typing import Any

from backend.database.user_store_protocol import IUserStore
from backend.models.events import EventKind
from dataikuapi.dss.llm_tracing import SpanBuilder

CONNECTION_CREDENTIALS_ACTION_KEY = "connection_credentials_issue"
CONNECTION_CREDENTIALS_ISSUE_FIELDS = {
    "connection_name",
    "connection_type",
    "message",
    "resolved",
}
ERROR_EVENT_KINDS = {
    EventKind.AGENT_ERROR.value,
    EventKind.BASE_LLM_ERROR.value,
}

_MISSING_CONNECTION_CREDENTIALS_PATTERN = re.compile(
    r"does not have credentials for connection ['\"]([^'\"]+)['\"](?:\s+to access\s+([^\"')\n\r.]+))?",
    re.IGNORECASE,
)


def deduplicate_connection_credentials_issues(issues: list[dict[str, str]]) -> list[dict[str, str]]:
    deduped: dict[str, dict[str, str]] = {}
    for issue in issues:
        connection_name = issue.get("connection_name", "").strip()
        if not connection_name:
            continue

        existing_issue = deduped.get(connection_name)
        if existing_issue is None:
            new_issue = {"connection_name": connection_name}
            for field, value in issue.items():
                if field in CONNECTION_CREDENTIALS_ISSUE_FIELDS and field != "connection_name":
                    new_issue[field] = value

            deduped[connection_name] = new_issue
            continue

        for field in ("message", "connection_type"):
            if not existing_issue.get(field) and issue.get(field):
                existing_issue[field] = issue[field]

    return list(deduped.values())


def _event_kind_value(event_kind: Any) -> Any:
    return getattr(event_kind, "value", event_kind)


def extract_connection_credentials_issues_from_error(error: Any) -> list[dict[str, str]]:
    """Return structured metadata for all missing-credentials matches in an error string."""
    if not isinstance(error, str):
        return []

    issues = []
    for match in _MISSING_CONNECTION_CREDENTIALS_PATTERN.finditer(error):
        issue = {
            "connection_name": match.group(1).strip(),
            "message": error,
        }
        connection_type = (match.group(2) or "").strip()
        if connection_type:
            issue["connection_type"] = connection_type
        issues.append(issue)

    return issues


def _extract_connection_credentials_issues_from_trace_span(span: dict[str, Any]) -> list[dict[str, str]]:
    """
    Extract issues from a SpanBuilder-style span.

    For typed spans, avoid walking all strings in inputs/attributes. Tool failures are written to
    outputs.error, while nested tool/LLM traces live in children.
    """
    issues: list[dict[str, str]] = []
    spans_to_visit = [span]

    while spans_to_visit:
        current_span = spans_to_visit.pop()

        outputs = current_span.get("outputs") or {}
        if isinstance(outputs, dict):
            issues.extend(extract_connection_credentials_issues_from_error(outputs.get("error")))

        children = current_span.get("children") or []
        if isinstance(children, list):
            spans_to_visit.extend(child for child in reversed(children) if isinstance(child, dict))

    return issues


def extract_connection_credentials_issues_from_trace(trace: SpanBuilder | None) -> list[dict[str, str]]:
    """Return all missing-credentials issues found in a trace payload."""
    if trace is None:
        return []

    return _extract_connection_credentials_issues_from_trace_span(trace.span)


def extract_connection_credentials_issues_from_events(events: Any) -> list[dict[str, str]]:
    """Return all missing-credentials issues found in streamed error events."""
    issues: list[dict[str, str]] = []
    for event in events or []:
        if not isinstance(event, dict):
            continue

        event_kind = _event_kind_value(event.get("eventKind"))
        if not isinstance(event_kind, str) or event_kind not in ERROR_EVENT_KINDS:
            continue

        event_data = event.get("eventData")
        if not isinstance(event_data, dict):
            continue

        issues.extend(extract_connection_credentials_issues_from_error(event_data.get("message")))

    return issues


def _normalize_connection_credentials_issue(issue: Any) -> dict | None:
    if not isinstance(issue, dict):
        return None

    return {
        "connection_name": issue.get("connection_name") or "",
        "connection_type": issue.get("connection_type") or "",
        "message": issue.get("message") or "",
        "resolved": issue.get("resolved") is True,
    }


def _get_connection_credentials_issues_from_event_data(data: Any) -> list[dict]:
    if not isinstance(data, dict):
        return []

    issues = []
    for issue in data.get("issues") or []:
        normalized_issue = _normalize_connection_credentials_issue(issue)
        if normalized_issue is not None:
            issues.append(normalized_issue)

    return issues


def get_connection_credentials_issue(events: list[dict]) -> dict | None:
    for event in events or []:
        if _event_kind_value(event.get("eventKind")) != EventKind.CONNECTION_CREDENTIALS_REQUIRED.value:
            continue

        issues = _get_connection_credentials_issues_from_event_data(event.get("eventData") or {})
        if issues:
            return {
                "issues": issues,
                "resolved": all(issue["resolved"] for issue in issues),
            }

    return None


def append_connection_credentials_issues(
    events: list[dict],
    issues: list[dict[str, str]],
) -> None:
    if not issues or get_connection_credentials_issue(events):
        return

    event_issues = []
    for issue in issues:
        event_issue = {key: value for key, value in issue.items() if key in CONNECTION_CREDENTIALS_ISSUE_FIELDS}
        event_issues.append(event_issue)

    event_issues = deduplicate_connection_credentials_issues(event_issues)
    if event_issues:
        events.append(
            {
                "eventKind": EventKind.CONNECTION_CREDENTIALS_REQUIRED,
                "eventData": {"issues": event_issues},
            }
        )


def append_connection_credentials_issue_from_response(events: list[dict], trace: SpanBuilder | None) -> None:
    if get_connection_credentials_issue(events):
        return

    issues: list[dict[str, str]] = []
    if trace:
        issues.extend(extract_connection_credentials_issues_from_trace(trace))
    issues.extend(extract_connection_credentials_issues_from_events(events))
    append_connection_credentials_issues(events, issues)


def mark_connection_credentials_issue_resolved(
    actions: dict,
    connection_name: str,
) -> dict | None:
    issue = actions.get(CONNECTION_CREDENTIALS_ACTION_KEY)
    issues = issue.get("issues") if isinstance(issue, dict) else None
    if not isinstance(issues, list) or not issues:
        return None

    updated_issues = []
    matched_issue = False
    for item in issues:
        if not isinstance(item, dict):
            continue

        updated_item = {**item}
        if item.get("connection_name") == connection_name:
            updated_item["resolved"] = True
            matched_issue = True
        updated_issues.append(updated_item)

    if not matched_issue:
        return None

    updated_issue = {
        "issues": updated_issues,
        "resolved": all(item.get("resolved") is True for item in updated_issues),
    }
    return {**actions, CONNECTION_CREDENTIALS_ACTION_KEY: updated_issue}


def resolve_connection_credentials_issue_for_message(
    store: IUserStore,
    message_id: str,
    connection_name: str,
) -> tuple[dict | None, str | None]:
    msg = store.get_message(message_id)
    if not msg:
        return None, "Message not found"

    actions = dict(msg.actions or {})
    updated_actions = mark_connection_credentials_issue_resolved(actions, connection_name)
    if not updated_actions:
        return None, "Connection credentials issue not found"

    store.update_message(message_id, {"actions": updated_actions})
    return updated_actions[CONNECTION_CREDENTIALS_ACTION_KEY], None


def resolve_connection_credentials_issue_for_draft_message(
    store: IUserStore,
    agent_id: str,
    message_id: str,
    connection_name: str,
) -> tuple[dict | None, str | None]:
    draft_conversation = store.get_draft_conversation(agent_id)
    if not draft_conversation or not draft_conversation.convo:
        return None, "Draft conversation not found"

    convo = draft_conversation.convo
    for message in convo.get("messages", []):
        if message.get("id") != message_id:
            continue

        actions = dict(message.get("actions") or {})
        updated_actions = mark_connection_credentials_issue_resolved(actions, connection_name)
        if not updated_actions:
            return None, "Connection credentials issue not found"

        message["actions"] = updated_actions
        store.upsert_draft_conversation(agent_id, convo)
        return updated_actions[CONNECTION_CREDENTIALS_ACTION_KEY], None

    return None, "Message not found"
