import base64
import re
from pathlib import Path
from urllib.parse import urlparse

from pydantic import BaseModel, Field

from ..typings import PydanticModelConfig, is_http_url

URL_PATTERN = re.compile(r"https?://[^\s<>\[\]\"'`]+", re.IGNORECASE)
URL_TRAILING_CHARS = ".,;:!?)]}>\"'"
IMAGE_EXTENSIONS = frozenset(
    {
        ".jpg",
        ".jpeg",
        ".png",
        ".gif",
        ".webp",
        ".bmp",
        ".svg",
        ".tif",
        ".tiff",
        ".heic",
        ".heif",
        ".avif",
    }
)


class PromptUrlExtractionResult(BaseModel):
    model_config = PydanticModelConfig.default()
    all_urls: list[str] = Field(default_factory=list)
    image_urls: list[str] = Field(default_factory=list)
    webpage_urls: list[str] = Field(default_factory=list)


def convert_image_to_base64(image_path: Path) -> str | None:
    """Convert an image file to a base64 encoded string."""
    if image_path is None:
        return None
    with open(image_path, "rb") as image_file:
        image_data = image_file.read()
        return base64.b64encode(image_data).decode("utf-8")


def is_image_url(url: str) -> bool:
    parsed = urlparse(url)
    path = parsed.path.lower()
    return any(path.endswith(extension) for extension in IMAGE_EXTENSIONS)


def extract_urls_from_prompt(prompt: str | None) -> PromptUrlExtractionResult:
    if prompt is None or prompt.strip() == "":
        return PromptUrlExtractionResult()

    ordered_urls: list[str] = []
    seen_urls: set[str] = set()

    for raw_url in URL_PATTERN.findall(prompt):
        normalized_url = raw_url.rstrip(URL_TRAILING_CHARS)
        if normalized_url == "":
            continue
        if not is_http_url(normalized_url):
            continue
        if normalized_url in seen_urls:
            continue
        seen_urls.add(normalized_url)
        ordered_urls.append(normalized_url)

    image_urls = [url for url in ordered_urls if is_image_url(url)]
    image_url_set = set(image_urls)
    webpage_urls = [url for url in ordered_urls if url not in image_url_set]

    return PromptUrlExtractionResult(
        all_urls=ordered_urls,
        image_urls=image_urls,
        webpage_urls=webpage_urls,
    )
