from __future__ import annotations

import json
import os
import re
import socket
import sys
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

DEFAULT_MAX_ATTEMPTS = 4

DEFAULT_BACKOFF_BASE = 1.0

DEFAULT_BACKOFF_MAX = 8.0


RETRYABLE_HTTP_STATUS_CODES = {
    408,
    425,
    429,
}


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


TOOL_ROUTING_SYSTEM_PROMPT = """
Si prvá fáza RAG agenta pre ZP Wiki.

Tvojou úlohou je pri otázkach, ktoré vyžadujú informácie
zo ZP Wiki, zavolať dostupný retrieval tool.

Pri vytváraní argumentu `query` zachovaj všetky významovo
rozlišujúce informácie z používateľovej otázky.

Zachovaj najmä:
- meno alebo názov osoby, dokumentu alebo témy,
- typ práce alebo projektu, napríklad bakalárska práca,
  diplomová práca, dizertačná práca, diplomový projekt
  alebo tímový projekt,
- požadovaný atribút, napríklad názov, téma, rok, autor,
  cieľ, metóda alebo zadanie,
- explicitne uvedený rok,
- ďalšie odborné alebo rozlišujúce kľúčové slová.

Query môže byť stručnejší než pôvodná otázka, ale nesmie
stratiť významové obmedzenia potrebné na správne vyhľadanie.

Napríklad otázku:
"Aká téma diplomovej práce je uvedená pri osobe Ján Novák?"
je vhodné previesť na query podobný:
"Ján Novák diplomová práca téma"

Nevhodné je zredukovať ju iba na:
"Ján Novák"

Podobne pri otázke na rok diplomovej práce musí query
zachovať meno osoby, pojem diplomová práca a informáciu,
že sa hľadá rok.

Nevymýšľaj hodnoty, ktoré používateľ v otázke neuviedol.
Nevkladaj do query predpokladanú odpoveď.

Ak je potrebné použiť ZP Wiki, najprv zavolaj retrieval tool.
Po získaní výsledku odpovedaj iba podľa informácií,
inštrukcií, kontextu a zdrojov vrátených toolom.
""".strip()


class RequestError(
    RuntimeError
):
    def __init__(
        self,
        message: str,
        *,
        retryable: bool,
        status_code: int | None = None,
    ) -> None:
        super().__init__(
            message
        )

        self.retryable = (
            retryable
        )

        self.status_code = (
            status_code
        )


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

    if value:
        return value

    if not allow_env_file:
        raise RuntimeError(
            f"{key} musí byť nastavený "
            "priamo v environment."
        )

    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 retry_delay_seconds(
    attempt: int,
    *,
    backoff_base: float,
    backoff_max: float,
) -> float:
    if attempt <= 0:
        return 0.0

    delay = (
        backoff_base
        * (
            2
            ** (
                attempt
                - 1
            )
        )
    )

    return min(
        delay,
        backoff_max,
    )


def request_json(
    url: str,
    *,
    method: str = "GET",
    headers: dict[str, str] | None = None,
    payload: dict[str, Any] | None = None,
    timeout: int = DEFAULT_TIMEOUT,
    max_attempts: int = 1,
    backoff_base: float = DEFAULT_BACKOFF_BASE,
    backoff_max: float = DEFAULT_BACKOFF_MAX,
) -> dict[str, Any]:
    if timeout <= 0:
        raise ValueError(
            "timeout musí byť > 0"
        )

    if max_attempts <= 0:
        raise ValueError(
            "max_attempts musí byť > 0"
        )

    if backoff_base < 0:
        raise ValueError(
            "backoff_base nesmie byť záporné"
        )

    if backoff_max < 0:
        raise ValueError(
            "backoff_max nesmie byť záporné"
        )

    data = None

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

    last_error: RequestError | None = None

    for attempt in range(
        1,
        max_attempts + 1,
    ):
        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",
                )
            )

            status_code = int(
                exc.code
            )

            retryable = (
                status_code
                in RETRYABLE_HTTP_STATUS_CODES
                or (
                    500
                    <= status_code
                    <= 599
                )
            )

            error = RequestError(
                (
                    f"HTTP {status_code} "
                    f"pre {url}: "
                    f"{body[:1500]}"
                ),
                retryable=(
                    retryable
                ),
                status_code=(
                    status_code
                ),
            )

        except (
            urllib.error.URLError,
            TimeoutError,
            socket.timeout,
        ) as exc:
            error = RequestError(
                (
                    "Sieťová chyba "
                    f"pre {url}: "
                    f"{exc}"
                ),
                retryable=True,
            )

        else:
            if not raw.strip():
                raise RequestError(
                    (
                        "Prázdna odpoveď "
                        f"z {url}."
                    ),
                    retryable=False,
                )

            try:
                parsed = json.loads(
                    raw
                )

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

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

            return parsed

        last_error = error

        if (
            not error.retryable
            or attempt >= max_attempts
        ):
            raise error

        delay = retry_delay_seconds(
            attempt,
            backoff_base=(
                backoff_base
            ),
            backoff_max=(
                backoff_max
            ),
        )

        print(
            (
                "    retry HTTP request: "
                f"pokus {attempt + 1}/"
                f"{max_attempts} "
                f"za {delay:.1f}s "
                f"({error})"
            ),
            file=sys.stderr,
        )

        if delay > 0:
            time.sleep(
                delay
            )

    if last_error is not None:
        raise last_error

    raise RuntimeError(
        "HTTP request skončil "
        "v neočakávanom stave."
    )


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 build_initial_messages(
    question: str,
) -> list[dict[str, Any]]:
    question = (
        question.strip()
    )

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

    return [
        {
            "role": "system",
            "content": (
                TOOL_ROUTING_SYSTEM_PROMPT
            ),
        },
        {
            "role": "user",
            "content": question,
        },
    ]


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,
    max_attempts: int = DEFAULT_MAX_ATTEMPTS,
    backoff_base: float = DEFAULT_BACKOFF_BASE,
    backoff_max: float = DEFAULT_BACKOFF_MAX,
) -> 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 = (
        build_initial_messages(
            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,
        max_attempts=max_attempts,
        backoff_base=backoff_base,
        backoff_max=backoff_max,
    )

    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,
            max_attempts=max_attempts,
            backoff_base=backoff_base,
            backoff_max=backoff_max,
        )

        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,
        max_attempts=max_attempts,
        backoff_base=backoff_base,
        backoff_max=backoff_max,
    )

    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"
            )
        ),
    }