from __future__ import annotations

import json
import os
import re
import time
import urllib.error
import urllib.request
from pathlib import Path
from typing import Any


PROJECT_ROOT = Path(
    __file__
).resolve().parents[1]


OPENWEBUI_URL = (
    "https://ui.tukekemt.xyz/api/chat/completions"
)

LOCAL_OPENAPI_URL = (
    "http://localhost:8000/openapi.json"
)

LOCAL_RAG_URL = (
    "http://localhost:8000/rag"
)

DEFAULT_MODEL = (
    "model120-fast"
)

DEFAULT_TIMEOUT = 180


MARKDOWN_URL_RE = re.compile(
    r"\[[^\]]*\]\((https?://[^)]+)\)"
)


def load_env_value(
    key: str,
    env_path: Path | None = None,
) -> str:
    value = os.environ.get(
        key
    )

    if value:
        return value

    if env_path is None:
        env_path = (
            PROJECT_ROOT
            / ".env"
        )

    if not env_path.exists():
        raise RuntimeError(
            f"{key} nie je v environment "
            f"a {env_path} neexistuje."
        )

    for raw_line in env_path.read_text(
        encoding="utf-8"
    ).splitlines():
        line = raw_line.strip()

        if (
            not line
            or line.startswith("#")
            or "=" not in line
        ):
            continue

        name, value = line.split(
            "=",
            1,
        )

        if name.strip() != key:
            continue

        value = value.strip()

        if (
            len(value) >= 2
            and value[0] == value[-1]
            and value[0] in {
                "'",
                '"',
            }
        ):
            value = value[
                1:-1
            ]

        if value:
            return value

    raise RuntimeError(
        f"{key} sa nepodarilo nájsť."
    )


def request_json(
    url: str,
    *,
    method: str = "GET",
    headers: dict[str, str] | None = None,
    payload: dict[str, Any] | None = None,
    timeout: int = DEFAULT_TIMEOUT,
) -> dict[str, Any]:
    data = None

    if payload is not None:
        data = json.dumps(
            payload,
            ensure_ascii=False,
        ).encode(
            "utf-8"
        )

    request = urllib.request.Request(
        url,
        data=data,
        method=method,
        headers=(
            headers
            or {}
        ),
    )

    try:
        with urllib.request.urlopen(
            request,
            timeout=timeout,
        ) as response:
            raw = (
                response
                .read()
                .decode(
                    "utf-8"
                )
            )

    except urllib.error.HTTPError as exc:
        body = (
            exc.read()
            .decode(
                "utf-8",
                errors="replace",
            )
        )

        raise RuntimeError(
            f"HTTP {exc.code} pre {url}: "
            f"{body[:1500]}"
        ) from exc

    except urllib.error.URLError as exc:
        raise RuntimeError(
            f"Sieťová chyba pre {url}: "
            f"{exc}"
        ) from exc

    if not raw.strip():
        raise RuntimeError(
            f"Prázdna odpoveď z {url}."
        )

    try:
        parsed = json.loads(
            raw
        )

    except json.JSONDecodeError as exc:
        raise RuntimeError(
            f"Neplatný JSON z {url}: "
            f"{raw[:1000]}"
        ) from exc

    if not isinstance(
        parsed,
        dict,
    ):
        raise RuntimeError(
            "Očakávaný JSON objekt "
            f"z {url}, dostal som "
            f"{type(parsed).__name__}."
        )

    return parsed


def resolve_refs(
    value: Any,
    document: dict[str, Any],
) -> Any:
    if isinstance(
        value,
        list,
    ):
        return [
            resolve_refs(
                item,
                document,
            )
            for item in value
        ]

    if not isinstance(
        value,
        dict,
    ):
        return value

    ref = value.get(
        "$ref"
    )

    if ref:
        if not ref.startswith(
            "#/"
        ):
            raise RuntimeError(
                "Nepodporovaný "
                f"OpenAPI $ref: {ref}"
            )

        current: Any = document

        for part in ref[
            2:
        ].split("/"):
            current = current[
                part
            ]

        resolved = resolve_refs(
            current,
            document,
        )

        extra = {
            key: item
            for key, item
            in value.items()
            if key != "$ref"
        }

        if (
            extra
            and isinstance(
                resolved,
                dict,
            )
        ):
            resolved = {
                **resolved,
                **resolve_refs(
                    extra,
                    document,
                ),
            }

        return resolved

    return {
        key: resolve_refs(
            item,
            document,
        )
        for key, item
        in value.items()
    }


def build_rag_tool(
    openapi: dict[str, Any],
) -> tuple[
    str,
    dict[str, Any],
]:
    try:
        operation = (
            openapi[
                "paths"
            ][
                "/rag"
            ][
                "post"
            ]
        )

    except KeyError as exc:
        raise RuntimeError(
            "OpenAPI schéma neobsahuje "
            "POST /rag."
        ) from exc

    operation_id = operation.get(
        "operationId"
    )

    if not operation_id:
        raise RuntimeError(
            "POST /rag nemá operationId."
        )

    try:
        schema = (
            operation[
                "requestBody"
            ][
                "content"
            ][
                "application/json"
            ][
                "schema"
            ]
        )

    except KeyError as exc:
        raise RuntimeError(
            "POST /rag nemá "
            "request JSON schema."
        ) from exc

    parameters = resolve_refs(
        schema,
        openapi,
    )

    description = (
        operation.get(
            "description"
        )
        or operation.get(
            "summary"
        )
        or (
            "Vyhľadá relevantný kontext "
            "v dokumentoch ZP Wiki."
        )
    )

    tool = {
        "type": "function",
        "function": {
            "name": operation_id,
            "description": (
                description
            ),
            "parameters": (
                parameters
            ),
        },
    }

    return (
        operation_id,
        tool,
    )


def normalize_url(
    value: str,
) -> str:
    value = value.strip()

    match = (
        MARKDOWN_URL_RE.search(
            value
        )
    )

    if match:
        value = match.group(
            1
        )

    return value.rstrip(
        "/"
    )


def extract_urls_from_object(
    value: Any,
) -> list[str]:
    result: list[
        str
    ] = []

    if isinstance(
        value,
        dict,
    ):
        for key, item in value.items():
            if (
                key == "source_url"
                and isinstance(
                    item,
                    str,
                )
            ):
                result.append(
                    normalize_url(
                        item
                    )
                )

            result.extend(
                extract_urls_from_object(
                    item
                )
            )

    elif isinstance(
        value,
        list,
    ):
        for item in value:
            result.extend(
                extract_urls_from_object(
                    item
                )
            )

    return list(
        dict.fromkeys(
            result
        )
    )


def get_usage(
    response: dict[str, Any],
) -> dict[str, int]:
    usage = response.get(
        "usage"
    )

    if not isinstance(
        usage,
        dict,
    ):
        return {
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "total_tokens": 0,
        }

    return {
        "prompt_tokens": int(
            usage.get(
                "prompt_tokens",
                0,
            )
            or 0
        ),
        "completion_tokens": int(
            usage.get(
                "completion_tokens",
                0,
            )
            or 0
        ),
        "total_tokens": int(
            usage.get(
                "total_tokens",
                0,
            )
            or 0
        ),
    }


def add_usage(
    total: dict[str, int],
    current: dict[str, int],
) -> None:
    for key in (
        "prompt_tokens",
        "completion_tokens",
        "total_tokens",
    ):
        total[
            key
        ] += current[
            key
        ]


def get_first_message(
    response: dict[str, Any],
) -> dict[str, Any]:
    choices = response.get(
        "choices"
    )

    if (
        not isinstance(
            choices,
            list,
        )
        or not choices
    ):
        raise RuntimeError(
            "Model nevrátil choices."
        )

    choice = choices[
        0
    ]

    if not isinstance(
        choice,
        dict,
    ):
        raise RuntimeError(
            "Neplatný choices[0]."
        )

    message = choice.get(
        "message"
    )

    if not isinstance(
        message,
        dict,
    ):
        raise RuntimeError(
            "Model nevrátil message."
        )

    return message


def parse_tool_arguments(
    raw_arguments: Any,
) -> dict[str, Any]:
    if isinstance(
        raw_arguments,
        dict,
    ):
        return raw_arguments

    if not isinstance(
        raw_arguments,
        str,
    ):
        raise RuntimeError(
            "Neplatný formát "
            "tool arguments."
        )

    try:
        parsed = json.loads(
            raw_arguments
        )

    except json.JSONDecodeError as exc:
        raise RuntimeError(
            "Model vrátil neplatné JSON "
            "argumenty toolu: "
            f"{raw_arguments[:1000]}"
        ) from exc

    if not isinstance(
        parsed,
        dict,
    ):
        raise RuntimeError(
            "Tool arguments nie sú "
            "JSON objekt."
        )

    return parsed


def run_question(
    question: str,
    *,
    model: str,
    operation_id: str,
    rag_tool: dict[str, Any],
    openwebui_api_key: str,
    search_api_key: str,
    timeout: int,
) -> dict[str, Any]:
    question = question.strip()

    if not question:
        raise RuntimeError(
            "Otázka je prázdna."
        )

    started = (
        time.perf_counter()
    )

    usage_total = {
        "prompt_tokens": 0,
        "completion_tokens": 0,
        "total_tokens": 0,
    }

    tool_calls_record: list[
        dict[str, Any]
    ] = []

    rag_source_urls: list[
        str
    ] = []

    messages: list[
        dict[str, Any]
    ] = [
        {
            "role": "user",
            "content": question,
        }
    ]

    first_started = (
        time.perf_counter()
    )

    first_response = request_json(
        OPENWEBUI_URL,
        method="POST",
        headers={
            "Authorization": (
                "Bearer "
                f"{openwebui_api_key}"
            ),
            "Content-Type": (
                "application/json"
            ),
        },
        payload={
            "model": model,
            "messages": messages,
            "tools": [
                rag_tool
            ],
            "tool_choice": "auto",
            "stream": False,
        },
        timeout=timeout,
    )

    first_latency = (
        time.perf_counter()
        - first_started
    )

    add_usage(
        usage_total,
        get_usage(
            first_response
        ),
    )

    first_message = (
        get_first_message(
            first_response
        )
    )

    tool_calls = (
        first_message.get(
            "tool_calls"
        )
        or []
    )

    if not isinstance(
        tool_calls,
        list,
    ):
        tool_calls = []

    if not tool_calls:
        answer = str(
            first_message.get(
                "content"
            )
            or ""
        )

        total_latency = (
            time.perf_counter()
            - started
        )

        return {
            "answer": answer,
            "tool_called": False,
            "tool_call_count": 0,
            "tool_calls": [],
            "rag_source_urls": [],
            "first_model_latency_seconds": (
                round(
                    first_latency,
                    6,
                )
            ),
            "tool_latency_seconds": 0.0,
            "final_model_latency_seconds": 0.0,
            "total_latency_seconds": (
                round(
                    total_latency,
                    6,
                )
            ),
            "usage": (
                usage_total
            ),
            "response_model": (
                first_response.get(
                    "model"
                )
            ),
        }

    messages.append(
        {
            "role": "assistant",
            "content": (
                first_message.get(
                    "content"
                )
                or ""
            ),
            "tool_calls": (
                tool_calls
            ),
        }
    )

    tool_latency_total = 0.0

    for tool_call in tool_calls:
        if not isinstance(
            tool_call,
            dict,
        ):
            raise RuntimeError(
                "Neplatný tool_call objekt."
            )

        function = tool_call.get(
            "function"
        )

        if not isinstance(
            function,
            dict,
        ):
            raise RuntimeError(
                "tool_call nemá function."
            )

        name = function.get(
            "name"
        )

        if name != operation_id:
            raise RuntimeError(
                "Model zavolal "
                "neočakávaný tool: "
                f"{name!r}"
            )

        arguments = (
            parse_tool_arguments(
                function.get(
                    "arguments"
                )
            )
        )

        tool_started = (
            time.perf_counter()
        )

        rag_result = request_json(
            LOCAL_RAG_URL,
            method="POST",
            headers={
                "X-API-Key": (
                    search_api_key
                ),
                "Content-Type": (
                    "application/json"
                ),
            },
            payload=arguments,
            timeout=timeout,
        )

        tool_latency = (
            time.perf_counter()
            - tool_started
        )

        tool_latency_total += (
            tool_latency
        )

        current_urls = (
            extract_urls_from_object(
                rag_result
            )
        )

        for url in current_urls:
            if (
                url
                not in rag_source_urls
            ):
                rag_source_urls.append(
                    url
                )

        tool_calls_record.append(
            {
                "name": name,
                "arguments": arguments,
                "latency_seconds": (
                    round(
                        tool_latency,
                        6,
                    )
                ),
                "source_urls": (
                    current_urls
                ),
            }
        )

        tool_call_id = (
            tool_call.get(
                "id"
            )
        )

        if not tool_call_id:
            raise RuntimeError(
                "tool_call nemá id."
            )

        messages.append(
            {
                "role": "tool",
                "tool_call_id": (
                    tool_call_id
                ),
                "name": name,
                "content": json.dumps(
                    rag_result,
                    ensure_ascii=False,
                ),
            }
        )

    final_started = (
        time.perf_counter()
    )

    final_response = request_json(
        OPENWEBUI_URL,
        method="POST",
        headers={
            "Authorization": (
                "Bearer "
                f"{openwebui_api_key}"
            ),
            "Content-Type": (
                "application/json"
            ),
        },
        payload={
            "model": model,
            "messages": messages,
            "stream": False,
        },
        timeout=timeout,
    )

    final_latency = (
        time.perf_counter()
        - final_started
    )

    add_usage(
        usage_total,
        get_usage(
            final_response
        ),
    )

    final_message = (
        get_first_message(
            final_response
        )
    )

    answer = str(
        final_message.get(
            "content"
        )
        or ""
    )

    total_latency = (
        time.perf_counter()
        - started
    )

    return {
        "answer": answer,
        "tool_called": True,
        "tool_call_count": len(
            tool_calls_record
        ),
        "tool_calls": (
            tool_calls_record
        ),
        "rag_source_urls": (
            rag_source_urls
        ),
        "first_model_latency_seconds": (
            round(
                first_latency,
                6,
            )
        ),
        "tool_latency_seconds": (
            round(
                tool_latency_total,
                6,
            )
        ),
        "final_model_latency_seconds": (
            round(
                final_latency,
                6,
            )
        ),
        "total_latency_seconds": (
            round(
                total_latency,
                6,
            )
        ),
        "usage": (
            usage_total
        ),
        "response_model": (
            final_response.get(
                "model"
            )
        ),
    }