import datetime
import json
import re
import time
from typing import Any, List
from dataiku.llm.tracing import new_trace

import dataiku

from backend.config import get_title_generation_llm
from backend.database.user_store_protocol import IUserStore
from backend.schemas import schemas
from backend.utils.logging_utils import get_logger
from backend.config import is_answers_mode_on

logger = get_logger(__name__)


def _clean_title(raw: str) -> str:
    s = (raw or "").strip()
    s = re.sub(r"```[\w\-]*\s*([\s\S]*?)\s*```", r"\1", s).strip()

    for line in s.splitlines():
        if line.strip():
            s = line.strip()
            break
    else:
        s = ""

    s = re.sub(r"^\s*#{1,6}\s*", "", s)  # '#', '##', etc.
    s = re.sub(r"^\s*>+\s*", "", s)  # '>' or '>>'

    s = re.sub(
        r"^\s*(?:title|subject|chat|conversation)\s*[:：\-–]\s*",
        "",
        s,
        flags=re.IGNORECASE,
    )

    s = s.strip(" \"'`*_")

    s = re.sub(r"\s+", " ", s)

    s = re.sub(r"[.:;!?，。；！]+$", "", s)

    return s or "New Chat"


def _generate_title(messages: List[schemas.MessageBase], trace=None) -> str:
    if trace:
        title_span = trace.subspan("generate_title")
        title_span.begin(int(time.time() * 1000))
    else:
        title_span = None

    today = datetime.date.today().strftime("%Y-%m-%d")
    llm_id = get_title_generation_llm()

    # Add inputs to trace
    if title_span:
        title_span.inputs["llm_id"] = llm_id
        title_span.inputs["message_count"] = len(messages)
        title_span.inputs["date"] = today

    system_msg = (
        "You are a titling assistant. "
        f"Today is {today}. "
        "Return a single short title (≤7 words) that summarizes the conversation. "
        "Plain text ONLY: no prefixes (e.g., 'Title:'), no quotes, no emojis, "
        "no code fences, no extra lines, and no trailing punctuation."
    )

    user_msg = "Generate a short title for this conversation. Output only the title text."

    project = dataiku.api_client().get_default_project()
    comp = project.get_llm(llm_id).new_completion()
    comp.with_message(system_msg, role="system")

    for m in messages:
        comp = comp.with_message(m.content, role=m.role)

    # the title generation instruction is moved at the end for
    # the recency bias and to support Mistral which doesn't like
    # to have assistant as last message role
    comp.with_message(user_msg, role="user")

    logger.info(
        "Generating conversation title:\nLLM id=[%s]\nCompletion Query=%s\nCompletion Settings=%s\n",
        llm_id,
        json.dumps(comp.cq, indent=2, sort_keys=True),
        json.dumps(comp.settings, indent=2, sort_keys=True),
    )

    resp = comp.execute()
    raw_title = resp.text or ""
    logger.info(
        "Conversation title LLM response:\nLLM id=[%s]\nCompletion Response=%s",
        llm_id,
        raw_title,
    )

    if title_span and resp.trace:
        title_span.append_trace(resp.trace)

    cleaned_title = _clean_title(raw_title)

    # Add outputs to trace
    if title_span:
        title_span.outputs["raw_title"] = raw_title
        title_span.outputs["cleaned_title"] = cleaned_title
        title_span.end(int(time.time() * 1000))
    return cleaned_title


def generate_title_with_trace(messages: List[schemas.MessageBase]) -> tuple[str | None, dict[str, Any]]:
    trace_title = "DKU_AGENT_CHAT_TITLE_GENERATION" if is_answers_mode_on() else "DKU_AGENT_HUB_TITLE_GENERATION"
    local_trace = new_trace(trace_title)
    local_trace.begin(int(time.time() * 1000))
    generated_title: str | None = None

    try:
        generated_title = _generate_title(messages, trace=local_trace)
    except Exception as e:
        logger.exception("Failed to generate conversation title")
        local_trace.outputs["error"] = str(e)
    finally:
        local_trace.end(int(time.time() * 1000))

    return generated_title, local_trace.to_dict()


def update_conversation_title(
    store: IUserStore,
    conv_id: str,
    generated_title: str,
    emit,
):
    try:
        store.update_conversation_meta(conv_id, title=generated_title)
        # Emit title update event so frontend can update the UI
        emit(
            "chat_title_updated",
            {
                "conversationId": conv_id,
                "title": generated_title,
            },
        )
    except Exception as e:
        # If title generation fails (e.g., broken LLM connection), keep existing title
        logger.exception(f"Failed to generate title for conversation {conv_id}: {e}")
        fallback_title = "New Chat"

        # Persist fallback title in database
        try:
            store.update_conversation_meta(
                conv_id,
                title=fallback_title,
            )
        except Exception:
            logger.exception("Failed to persist fallback title for conversation %s", conv_id)
