import ast
import asyncio
import json
import re
from collections.abc import Coroutine, Mapping
from datetime import datetime
from logging import getLogger
from pathlib import Path
from typing import Any, TypeVar

from pydantic import (
    GetCoreSchemaHandler,
    TypeAdapter,
    ValidationError,
)
from pydantic_core import core_schema
from tqdm.autonotebook import tqdm

logger = getLogger(__name__)

_JSON_START_RE = re.compile(r"[{\[]")

T = TypeVar("T")


def extract_json_substring(text: str) -> str | None:
    decoder = json.JSONDecoder()
    for match in _JSON_START_RE.finditer(text):
        start = match.start()
        try:
            _, end = decoder.raw_decode(text, idx=start)
            return text[start:end]
        except ValueError:
            continue

    return None


def parse_json_or_py_string(
    s: str, return_none_on_failure: bool = False
) -> dict[str, Any] | list[Any] | None:
    s_fmt = re.sub(r"```[a-zA-Z0-9]*\n|```", "", s).strip()
    try:
        return ast.literal_eval(s_fmt)
    except (ValueError, SyntaxError):
        try:
            return json.loads(s_fmt)
        except json.JSONDecodeError as exc:
            if return_none_on_failure:
                return None
            raise ValueError(
                "Invalid JSON/Python string - Both ast.literal_eval and json.loads "
                f"failed to parse the following response:\n{s}"
            ) from exc


def parse_json_or_py_substring(
    json_str: str, return_none_on_failure: bool = False
) -> dict[str, Any] | list[Any] | None:
    return parse_json_or_py_string(
        extract_json_substring(json_str) or "", return_none_on_failure
    )


def validate_obj_from_json_or_py_string(
    s: str, adapter: TypeAdapter[T], from_substring: bool = False
) -> T:
    try:
        if from_substring:
            parsed = parse_json_or_py_substring(s, return_none_on_failure=True)
        else:
            parsed = parse_json_or_py_string(s, return_none_on_failure=True)
        if parsed is None:
            parsed = s
        return adapter.validate_python(parsed)
    except (json.JSONDecodeError, ValidationError) as exc:
        raise ValueError(
            f"Invalid JSON or Python string:\n{s}\nExpected type: {adapter._type}",  # type: ignore[arg-type]
        ) from exc


def extract_xml_list(text: str) -> list[str]:
    pattern = re.compile(r"<(chunk_\d+)>(.*?)</\1>", re.DOTALL)

    chunks: list[str] = []
    for match in pattern.finditer(text):
        content = match.group(2).strip()
        chunks.append(content)
    return chunks


def build_marker_json_parser_type(
    marker_to_model: Mapping[str, type],
) -> type:
    """
    Return a Pydantic-compatible *type* that, when given a **str**, searches for
    the first marker substring and validates the JSON that follows with the
    corresponding Pydantic model.

    If no marker is found, the raw string is returned unchanged.

    Example:
    -------
    >>> Todo = build_marker_json_parser_type({'```json': MyModel})
    >>> Todo.validate('```json {"a": 1}')
    MyModel(a=1)

    """

    class MarkerParsedOutput:
        """String → (Model | str) parser generated by build_marker_json_parser_type."""

        @classmethod
        def __get_pydantic_core_schema__(
            cls,
            _source_type: Any,
            _handler: GetCoreSchemaHandler,
        ) -> core_schema.CoreSchema:
            def _validate(value: Any) -> Any:
                if not isinstance(value, str):
                    raise TypeError("MarkerParsedOutput expects a string")

                for marker, model in marker_to_model.items():
                    if marker in value:
                        adapter = TypeAdapter[Any](model)
                        return validate_obj_from_json_or_py_string(
                            value, adapter=adapter, from_substring=True
                        )

                return value

            return core_schema.no_info_after_validator_function(
                _validate, core_schema.any_schema()
            )

        @classmethod
        def __get_pydantic_json_schema__(
            cls,
            schema: core_schema.CoreSchema,
            handler: GetCoreSchemaHandler,
        ):
            return handler(schema)

    unique_suffix = "_".join(sorted(marker_to_model))[:40]
    MarkerParsedOutput.__name__ = f"MarkerParsedOutput_{unique_suffix}"

    return MarkerParsedOutput


def read_txt(file_path: str | Path, encoding: str = "utf-8") -> str:
    return Path(file_path).read_text(encoding=encoding)


def read_contents_from_file(
    file_path: str | Path,
    binary_mode: bool = False,
) -> str | bytes:
    try:
        if binary_mode:
            return Path(file_path).read_bytes()
        return Path(file_path).read_text()
    except FileNotFoundError:
        logger.error(f"File {file_path} not found.")
        return ""


def get_prompt(prompt_text: str | None, prompt_path: str | Path | None) -> str | None:
    if prompt_text is None:
        return read_contents_from_file(prompt_path) if prompt_path is not None else None  # type: ignore[arg-type]

    return prompt_text


async def asyncio_gather_with_pbar(
    *corouts: Coroutine[Any, Any, Any],
    no_tqdm: bool = False,
    desc: str | None = None,
) -> list[Any]:
    pbar = tqdm(total=len(corouts), desc=desc, disable=no_tqdm)

    async def run_and_update(coro: Coroutine[Any, Any, Any]) -> Any:
        result = await coro
        pbar.update(1)
        return result

    wrapped_tasks = [run_and_update(c) for c in corouts]
    results = await asyncio.gather(*wrapped_tasks)
    pbar.close()

    return results


def get_timestamp() -> str:
    return datetime.now().strftime("%Y%m%d_%H%M%S")
