"use client";

import { useEffect, useRef, useState, useMemo } from "react";
import { useRouter } from "next/navigation";
import { useChat } from "ai/react";
import { Send, User, Bot } from "lucide-react";
import DashboardLayout from "@/components/layout/dashboard-layout";
import { api, ApiError } from "@/lib/api";
import { useToast } from "@/components/ui/use-toast";
import { Answer } from "@/components/chat/answer";

interface Message {
  id: string;
  role: "assistant" | "user" | "system" | "data";
  content: string;
  citations?: Citation[];
}

interface ChatMessage {
  id: number;
  content: string;
  role: "assistant" | "user";
  created_at: string;
}

interface Chat {
  id: number;
  title: string;
  messages: ChatMessage[];
}

interface Citation {
  id: number;
  text: string;
  metadata: Record<string, any>;
}

interface ContextPayload {
  context?: Array<{
    page_content: string;
    metadata: Record<string, any>;
  }>;
  route?: Record<string, any>;
  intent?: string;
  selected_chain?: string;
  retrieval_preview?: Array<Record<string, any>>;
}

// Extend the default useChat message type
declare module "ai/react" {
  interface Message {
    citations?: Citation[];
  }
}

const LLM_RESPONSE_SEPARATOR = "__LLM_RESPONSE__";
const LEADING_BASE64_RE = /^[A-Za-z0-9+/=]+/;
const CODE_FENCE_RE = /(```[\s\S]*?```)/g;
const ANSWER_KEYWORDS = [
  "测试充分性要求",
  "人机交互界面测试",
  "外部接口测试",
  "测试用例",
  "测试项",
  "预期成果",
  "正常测试",
  "异常测试",
  "总体结论",
  "核心结论",
  "关键结论",
  "关键点",
  "关键词",
  "注意事项",
  "步骤",
  "建议",
  "原因",
  "风险",
  "结论",
];

const escapeRegExp = (value: string): string =>
  value.replace(/[.*+?^${}()|[\]\\]/g, "\\$&");

const KEYWORD_PATTERN = new RegExp(
  ANSWER_KEYWORDS.map((keyword) => escapeRegExp(keyword))
    .sort((left, right) => right.length - left.length)
    .join("|"),
  "g"
);

const normalizeCitationSyntax = (text: string, maxCitationId: number): string => {
  const normalized = text
    .replace(/\[\[([cC])itation/g, "[citation")
    .replace(/[cC]itation:(\d+)]]/g, "citation:$1]")
    .replace(/\[\[([cC]itation:\d+)]](?!])/g, "[$1]")
    .replace(/\[[cC]itation:(\d+)]/g, "[citation]($1)");

  return normalized.replace(/\[(\d+)]/g, (raw, idText) => {
    const id = Number(idText);
    if (!Number.isFinite(id) || id <= 0 || id > maxCitationId) {
      return raw;
    }
    return `[citation](${id})`;
  });
};

const extractCitationIds = (text: string, maxCitationId: number): number[] => {
  const ids: number[] = [];
  const seen = new Set<number>();
  const pattern = /\[citation]\((\d+)\)/gi;

  let matched: RegExpExecArray | null = pattern.exec(text);
  while (matched) {
    const id = Number(matched[1]);
    if (id > 0 && id <= maxCitationId && !seen.has(id)) {
      seen.add(id);
      ids.push(id);
    }
    matched = pattern.exec(text);
  }

  return ids;
};

const stripCitationMarkers = (text: string): string =>
  text
    .replace(/\s*\[citation]\((\d+)\)\s*/gi, " ")
    .replace(/[ \t]+\n/g, "\n")
    .replace(/\n{3,}/g, "\n\n")
    .trim();

const highlightAnswerKeywords = (text: string): string => {
  if (!text.trim()) {
    return text;
  }

  return text
    .split(CODE_FENCE_RE)
    .map((segment) => {
      if (segment.startsWith("```")) {
        return segment;
      }

      const protectedBlocks: string[] = [];
      const protectedSegment = segment.replace(/\*\*[^*]+\*\*/g, (matched) => {
        const token = `__BOLD_BLOCK_${protectedBlocks.length}__`;
        protectedBlocks.push(matched);
        return token;
      });

      const highlighted = protectedSegment.replace(KEYWORD_PATTERN, (matched) => `**${matched}**`);

      return highlighted.replace(/__BOLD_BLOCK_(\d+)__/g, (_raw, indexText) => {
        const index = Number(indexText);
        return protectedBlocks[index] ?? "";
      });
    })
    .join("");
};

const formatAssistantMarkdown = (text: string, citations: Citation[]): string => {
  const maxCitationId = citations.length;
  const normalized = normalizeCitationSyntax(text, maxCitationId);
  const citationIds = extractCitationIds(normalized, maxCitationId);
  const fallbackCitationIds = citationIds.length
    ? citationIds
    : citations.map((citation) => citation.id);

  const body = stripCitationMarkers(normalized);
  const highlightedBody = highlightAnswerKeywords(body);

  if (!fallbackCitationIds.length) {
    return highlightedBody;
  }

  const refs = fallbackCitationIds.map((id) => `[citation](${id})`).join(" ");
  return `${highlightedBody}\n\n**引用**：${refs}`;
};

const decodeBase64Json = (encoded: string): ContextPayload | null => {
  try {
    const binary = atob(encoded);
    const bytes = Uint8Array.from(binary, (ch) => ch.charCodeAt(0));
    const decoded = new TextDecoder("utf-8").decode(bytes);
    return JSON.parse(decoded) as ContextPayload;
  } catch {
    return null;
  }
};

const isLikelyContextPayload = (payload: ContextPayload | null): boolean => {
  if (!payload || typeof payload !== "object") {
    return false;
  }

  return Boolean(
    payload.context ||
      payload.route ||
      payload.intent ||
      payload.selected_chain ||
      payload.retrieval_preview
  );
};

const parseAssistantPayload = (rawContent: string): {
  content: string;
  citations: Citation[];
} => {
  let contextPayload: ContextPayload | null = null;
  let responseText = rawContent;

  const separatorIndex = rawContent.indexOf(LLM_RESPONSE_SEPARATOR);

  if (separatorIndex >= 0) {
    const base64Part = rawContent.slice(0, separatorIndex).trim();
    contextPayload = decodeBase64Json(base64Part);
    responseText = rawContent.slice(separatorIndex + LLM_RESPONSE_SEPARATOR.length);
  } else {
    const base64Match = rawContent.match(LEADING_BASE64_RE);
    if (base64Match && base64Match[0].length > 80) {
      const candidate = base64Match[0];
      const candidatePayload = decodeBase64Json(candidate);
      if (isLikelyContextPayload(candidatePayload)) {
        contextPayload = candidatePayload;
        responseText = rawContent.slice(candidate.length);
      }
    }
  }

  responseText = responseText.replace(/^_+LLM_RESPONSE_+/, "").trimStart();

  const citations: Citation[] =
    contextPayload?.context?.map((citation, index) => ({
      id: index + 1,
      text: citation.page_content,
      metadata: citation.metadata,
    })) || [];

  return { content: responseText, citations };
};

export default function ChatPage({ params }: { params: { id: string } }) {
  const router = useRouter();
  const messagesEndRef = useRef<HTMLDivElement>(null);
  const { toast } = useToast();
  const [isInitialLoad, setIsInitialLoad] = useState(true);

  const {
    messages,
    data,
    input,
    handleInputChange,
    handleSubmit,
    isLoading,
    setMessages,
  } = useChat({
    api: `/api/chat/${params.id}/messages`,
    headers: {
      Authorization: `Bearer ${
        typeof window !== "undefined"
          ? window.localStorage.getItem("token")
          : ""
      }`,
    },
  });

  useEffect(() => {
    if (isInitialLoad) {
      fetchChat();
      setIsInitialLoad(false);
    }
  }, [isInitialLoad]);

  useEffect(() => {
    if (!isInitialLoad) {
      scrollToBottom();
    }
  }, [messages, isInitialLoad]);

  const fetchChat = async () => {
    try {
      const data: Chat = await api.get(`/api/chat/${params.id}`);
      const formattedMessages = data.messages.map((msg) => {
        if (msg.role !== "assistant" || !msg.content)
          return {
            id: msg.id.toString(),
            role: msg.role,
            content: msg.content,
          };

        try {
          const parsed = parseAssistantPayload(msg.content);

          return {
            id: msg.id.toString(),
            role: msg.role,
            content: parsed.content,
            citations: parsed.citations,
          };
        } catch (e) {
          console.error("Failed to process message:", e);
          return {
            id: msg.id.toString(),
            role: msg.role,
            content: msg.content,
          };
        }
      });
      setMessages(formattedMessages);
    } catch (error) {
      console.error("Failed to fetch chat:", error);
      if (error instanceof ApiError) {
        toast({
          title: "错误",
          description: error.message,
          variant: "destructive",
        });
      }
      router.push("/dashboard/chat");
    }
  };

  const scrollToBottom = () => {
    messagesEndRef.current?.scrollIntoView({ behavior: "smooth" });
  };

  const processedMessages = useMemo(() => {
    return messages.map((message) => {
      if (message.role !== "assistant" || !message.content) return message;

      try {
        const parsed = parseAssistantPayload(message.content);

        return {
          ...message,
          content: formatAssistantMarkdown(parsed.content, parsed.citations),
          citations: parsed.citations,
        };
      } catch (e) {
        console.error("Failed to process message:", e);
        return message;
      }
    });
  }, [messages]);

  return (
    <DashboardLayout>
      <div className="flex flex-col h-[calc(100vh-5rem)] relative">
        <div className="flex-1 overflow-y-auto p-4 space-y-4 pb-[80px]">
          {processedMessages.map((message) =>
            message.role === "assistant" ? (
              <div
                key={message.id}
                className="flex justify-start items-start space-x-2"
              >
                <div className="w-8 h-8 flex items-center justify-center">
                  <img
                    src="/logo.png"
                    className="h-8 w-8 rounded-full"
                    alt="logo"
                  />
                </div>
                <div className="max-w-[80%] rounded-lg px-4 py-2 text-accent-foreground">
                  <Answer
                    key={message.id}
                    markdown={message.content}
                    citations={message.citations}
                  />
                </div>
              </div>
            ) : (
              <div
                key={message.id}
                className="flex justify-end items-start space-x-2"
              >
                <div className="max-w-[80%] rounded-lg px-4 py-2 bg-primary text-primary-foreground">
                  {message.content}
                </div>
                <div className="w-8 h-8 rounded-full bg-primary flex items-center justify-center">
                  <User className="h-5 w-5 text-primary-foreground" />
                </div>
              </div>
            )
          )}
          <div className="flex justify-start">
            {isLoading &&
              processedMessages[processedMessages.length - 1]?.role !=
                "assistant" && (
                <div className="max-w-[80%] rounded-lg px-4 py-2 text-accent-foreground">
                  <div className="flex items-center space-x-1">
                    <div className="w-2 h-2 rounded-full bg-primary animate-bounce" />
                    <div className="w-2 h-2 rounded-full bg-primary animate-bounce [animation-delay:0.2s]" />
                    <div className="w-2 h-2 rounded-full bg-primary animate-bounce [animation-delay:0.4s]" />
                  </div>
                </div>
              )}
          </div>
          <div ref={messagesEndRef} />
        </div>
        <form
          onSubmit={handleSubmit}
          className="border-t p-4 flex items-center space-x-4 bg-background absolute bottom-0 left-0 right-0"
        >
          <input
            value={input}
            onChange={handleInputChange}
            placeholder="请输入消息..."
            className="flex-1 min-w-0 h-10 rounded-md border border-input bg-background px-3 py-2 text-sm ring-offset-background placeholder:text-muted-foreground focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring focus-visible:ring-offset-2"
          />
          <button
            type="submit"
            disabled={isLoading || !input.trim()}
            className="inline-flex items-center justify-center rounded-md text-sm font-medium ring-offset-background transition-colors focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring focus-visible:ring-offset-2 disabled:pointer-events-none disabled:opacity-50 bg-primary text-primary-foreground hover:bg-primary/90 h-10 px-4 py-2"
          >
            <Send className="h-4 w-4" />
          </button>
        </form>
      </div>
    </DashboardLayout>
  );
}
