from __future__ import annotations

from tools.evaluation.urls import MARKDOWN_URL_RE, normalize_url
from tools.rag.llm_context import compact_rag_payload
from tools.answers.rag_evidence import rag_prose_issues

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

from tools.answers.context_policy import citation_check, coverage_issues

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

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


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 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] = []
    retrieval_coverage = []
    delivery_contexts = []

    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

        rag_result = compact_rag_payload(rag_result)
        delivery_contexts.append(rag_result)
        current_urls = [
            source["source_url"]
            for source in rag_result.get("sources", [])
            if source.get("source_url")
        ]
        retrieval_coverage.append(
            {
                key: rag_result.get(key)
                for key in ["complete", "truncated", "truncation_reasons", "index_version"]
            }
        )

        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,
                    separators=(",", ":"),
                ),
            }
        )

    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

    citations = citation_check({"answer": answer, "source_urls": []}, set(rag_source_urls))
    policy_issues = sorted(
        {
            issue
            for coverage in retrieval_coverage
            for issue in coverage_issues({"answer": answer, "status": "ok"}, coverage)
        }
    )
    for context in delivery_contexts:
        policy_issues.extend(rag_prose_issues(question, answer, context))
    policy_issues = sorted(set(policy_issues))
    return {
        **citations,
        "retrieval_coverage": retrieval_coverage,
        "context_issues": policy_issues,
        "accepted_for_delivery": citations["citations_ok"] and not policy_issues,
        "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")),
    }