from copy import deepcopy
from typing import AsyncGenerator
from pydantic import BaseModel
import os
import json
import aiohttp
from dotenv import load_dotenv
import re

load_dotenv()


class PerplexitySearchEvent(BaseModel):
    chunk: str


PERPLEXITY_API_URL = "https://api.perplexity.ai/chat/completions"


def should_include_citation(citation: str, domain_filter: list[str]) -> bool:
    # 如果沒有任何過濾規則，接受所有網站
    if not domain_filter:
        return True

    # 分離排除規則和包含規則
    exclude_rules = [
        rule[1:].replace("*.", "") for rule in domain_filter if rule.startswith("-")
    ]
    include_rules = [
        rule.replace("*.", "") for rule in domain_filter if not rule.startswith("-")
    ]

    # 檢查是否符合任何排除規則
    for pattern in exclude_rules:
        if pattern in citation:
            return False

    # 如果沒有包含規則，且通過了排除規則檢查，就接受該網站
    if not include_rules:
        return True

    # 如果有包含規則，必須符合至少一個
    for pattern in include_rules:
        if pattern in citation:
            return True

    return False


def is_valid_domain(domain: str) -> bool:
    if not domain or "*." in domain:
        return False

    # 只允許包含 ://、.、% 和英數字的網址
    # ^ 表示開頭，$ 表示結尾
    # [a-zA-Z0-9] 表示英數字
    # [\\.\\:\\/\\%] 表示允許的特殊字符
    pattern = r"^[a-zA-Z0-9\\.\\:\\/\\%]+$"

    return bool(re.match(pattern, domain))


async def respond_with_perplexity_search(
    input_content,
    user_prompt_prefix,
    messages_for_llm,
    domain_filter: list[str],
    stream: bool = False,
    model: str = "sonar-reasoning-pro",
) -> AsyncGenerator[PerplexitySearchEvent, None]:
    if model not in ["sonar-reasoning-pro", "sonar-pro"]:
        model = "sonar-reasoning-pro"
    api_key = os.getenv("PPLX_API_KEY")
    if not api_key:
        raise ValueError("PPLX_API_KEY environment variable not set")

    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }
    messages = deepcopy(messages_for_llm)
    if len(messages) > 0 and messages[-1]["role"] == "user":
        messages.pop()
    if user_prompt_prefix:
        xml_input_content = f"<使用者提問>{input_content}</使用者提問>"
        messages.append(
            {"role": "user", "content": user_prompt_prefix + "\n\n" + xml_input_content}
        )
    else:
        messages.append({"role": "user", "content": input_content})
    filtered_domain_filter = []

    for domain in domain_filter:
        if domain and is_valid_domain(domain):
            filtered_domain_filter.append(domain)

    payload = {
        "model": model,
        "messages": messages,
        "temperature": 0.5,
        "stream": stream,
        "search_domain_filter": filtered_domain_filter,
    }
    try:
        async with aiohttp.ClientSession() as session:
            async with session.post(
                PERPLEXITY_API_URL, headers=headers, json=payload
            ) as response:
                if response.status != 200:
                    error_text = await response.text()
                    raise ValueError(f"Perplexity API error: {error_text}")

                if not stream:
                    # 非串流模式的處理
                    response_data = await response.json()
                    content = response_data["choices"][0]["message"]["content"]
                    yield PerplexitySearchEvent(chunk=content)

                    # 處理引用
                    citations = response_data.get("citations", [])
                    final_citations = [
                        citation
                        for citation in citations
                        if should_include_citation(citation, domain_filter)
                    ]

                    if final_citations:
                        references = f"\n\n參考來源：\n"
                        for citation in final_citations:
                            references += f"- [{citation}]({citation})\n"
                        yield PerplexitySearchEvent(chunk=references)
                    return

                # 串流模式的處理
                full_response = ""
                final_citations = []
                async for line in response.content:
                    if line:
                        line = line.decode("utf-8").strip()
                        if line.startswith("data: "):
                            line = line[6:]  # Remove 'data: ' prefix
                            if line == "[DONE]":
                                break

                            try:
                                chunk_data = json.loads(line)
                                # print(chunk_data)
                                if (
                                    chunk_data["choices"][0]
                                    .get("delta", {})
                                    .get("content")
                                ):
                                    content = chunk_data["choices"][0]["delta"][
                                        "content"
                                    ]
                                    full_response += content
                                    yield PerplexitySearchEvent(
                                        chunk=content,
                                    )
                                if not final_citations and chunk_data.get(
                                    "citations", []
                                ):
                                    # 發現 perplexity 不會都有 finish_reason 為 stop 的狀況，但是 citations 會有
                                    # if chunk_data["choices"][0]["finish_reason"] == "stop":
                                    citations = chunk_data.get("citations", [])
                                    final_citations = [
                                        citation
                                        for citation in citations
                                        if should_include_citation(
                                            citation, domain_filter
                                        )
                                    ]

                            except json.JSONDecodeError:
                                continue

                # 只在有符合條件的 citations 時才產生參考文獻
                if final_citations:
                    references = f"\n\n參考來源：\n"
                    for citation in final_citations:
                        references += f"- [{citation}]({citation})\n"
                    yield PerplexitySearchEvent(chunk=references)

    except Exception as e:
        import traceback

        traceback.print_exc()
        print(e)

    # answer_message = await cl.Message(content="").send()
    # full_response = ""
    # for response in responses:
    #     if response.candidates[0].finish_reason != Candidate.FinishReason.STOP:
    #         # await answer_message.stream_token(response.text)
    #         yield GeminiGroundingEvent(chunk=response.text)
    #         full_response += response.text
    #         if response.candidates[0].grounding_metadata:
    #             if len(response.candidates[0].grounding_metadata.grounding_chunks) > 0:
    #                 references = f"\n\n{tr('Sources:')}\n"
    #                 for grounding_chunk in response.candidates[
    #                     0
    #                 ].grounding_metadata.grounding_chunks:
    #                     references += f"- [{grounding_chunk.web.title}]({grounding_chunk.web.uri})\n"
    #                 # await answer_message.stream_token(references)
    #                 yield GeminiGroundingEvent(chunk=references)
    #     else:
    #         if response.candidates[0].grounding_metadata:
    #             if len(response.candidates[0].grounding_metadata.grounding_chunks) > 0:
    #                 references = f"\n\n{tr('Sources:')}\n"
    #                 for grounding_chunk in response.candidates[
    #                     0
    #                 ].grounding_metadata.grounding_chunks:
    #                     references += f"- [{grounding_chunk.web.title}]({grounding_chunk.web.uri})\n"
    #                 # await answer_message.stream_token(references)
    #                 yield GeminiGroundingEvent(chunk=references)
