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 evaluačný klient pre ZP Agent. Pri každej otázke najprv použi "
    "nástroj retrieve_zpwiki_context na vyhľadanie dôkazov v ZP Wiki. "
    "Neodpovedaj z vlastnej pamäte a pred prvým vyhľadaním nežiadaj "
    "používateľa o spresnenie. Pri tvorbe argumentu query zachovaj všetky významovo "
    "dôležité obmedzenia pôvodnej otázky a zachovaj celý význam pôvodnej otázky: "
    "meno osoby alebo názov dokumentu, typ práce "
    "(bakalárska, diplomová, diplomový projekt, stáž), požadovaný atribút "
    "(názov, rok, téma, hardvér, databáza, backend, počet, nástroj), rok "
    "a dôležité odborné kľúčové slová. Nezredukuj otázku iba na meno osoby. "
    "Ak je otázka krátka alebo kontextová, napríklad sa pýta na chatbot, "
    "Question Answering, anotáciu alebo projekt bez mena osoby, aj tak najprv "
    "vyhľadaj pôvodnú otázku alebo jej hlavné pojmy v ZP Wiki. Pri otázke na "
    "viacero dokumentov zachovaj v query všetky tematické pojmy. Po získaní "
    "výsledku odpovedz iba podľa vráteného kontextu a cituj iba source_url, "
    "ktoré tvoju odpoveď podporujú."
)


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 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} pre {url}: {body[:1500]}",
                retryable=retryable,
                status_code=status_code,
            )

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

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

            try:
                parsed = json.loads(raw)

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

            if not isinstance(parsed, dict):
                raise RequestError(
                    "Očakávaný JSON objekt "
                    f"z {url}, dostal som {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}/{max_attempts} "
            f"za {delay:.1f}s ({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(
                f"Nepodporovaný 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 forced_tool_choice(
    operation_id: str,
) -> dict[str, Any]:
    return {
        "type": "function",
        "function": {
            "name": operation_id,
        },
    }


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": (
                f"Bearer {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:
        tool_calls = [
            {
                "id": (
                    "fallback_zpwiki_search"
                ),
                "type": "function",
                "function": {
                    "name": (
                        operation_id
                    ),
                    "arguments": json.dumps(
                        {
                            "query": question,
                        },
                        ensure_ascii=False,
                    ),
                },
            }
        ]

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

        # Query vždy prepíšeme celou
        # pôvodnou používateľskou otázkou.
        # Model si môže ponechať iba
        # neškodné voliteľné parametre
        # ako limit/published_only/
        # max_per_document.
        arguments[
            "query"
        ] = question

        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": (
                f"Bearer {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"
            )
        ),
    }