from __future__ import annotations

import asyncio
import inspect
import json
import logging
import os
import re
import threading
from copy import deepcopy
from typing import Any, Callable, Dict, List, Optional, Sequence

from ..common.tree_tools import (
    build_parent_lookup,
    canonical_node_id,
    flatten_with_parent,
    normalize_navexa_tree_document,
    split_selected_with_ancestor_dedupe,
    strip_fields,
)
from ..common.utils import ChatGPT_API, extract_json, get_usage_summary
from ..llm.client import has_llm_credentials, resolve_model_name
from ..prompts.tree_rag import render_answer_prompt, render_tree_search_prompt
from .models import AnswerResult, ContextBundle, RAGResult, TreeReasoningResult
from .tree_io import fetch_document_tree

VERBOSITY_LOW = "low"
VERBOSITY_MEDIUM = "medium"
VERBOSITY_HIGH = "high"


def _normalize_verbosity(value: Optional[str]) -> str:
    raw = (value or VERBOSITY_MEDIUM).strip().lower()
    if raw in {"1", "low"}:
        return VERBOSITY_LOW
    if raw in {"3", "high", "debug", "detailed"}:
        return VERBOSITY_HIGH
    return VERBOSITY_MEDIUM


def _configure_logger(verbosity: str) -> logging.Logger:
    logger = logging.getLogger("navexa.reasoning")
    logger.setLevel(logging.DEBUG if verbosity == VERBOSITY_HIGH else logging.INFO)
    logger.propagate = False

    if not logger.handlers:
        handler = logging.StreamHandler()
        handler.setFormatter(logging.Formatter("[NAVEXA][%(levelname)s] %(message)s"))
        logger.addHandler(handler)

    for handler in logger.handlers:
        handler.setLevel(logging.DEBUG if verbosity == VERBOSITY_HIGH else logging.INFO)
    return logger


def _has_llm_credentials() -> bool:
    return has_llm_credentials()


def _resolve_model(user_model: Optional[str]) -> Optional[str]:
    return resolve_model_name(user_model, use_llm_requested=True)


def _coerce_tree_input(tree: Any) -> Dict[str, Any]:
    if isinstance(tree, dict) and "structure" in tree:
        return normalize_navexa_tree_document(tree)
    if isinstance(tree, list):
        return normalize_navexa_tree_document({"structure": tree})
    if isinstance(tree, dict):
        return normalize_navexa_tree_document({"structure": [tree]})
    raise ValueError("tree must be a dict/list Navexa tree object")


def _usage_delta(before: Dict[str, Any], after: Dict[str, Any]) -> Dict[str, Any]:
    keys = ["calls", "input_tokens", "cached_input_tokens", "output_tokens", "total_tokens", "estimated_cost_usd"]
    delta: Dict[str, Any] = {}
    for key in keys:
        if key == "estimated_cost_usd":
            delta[key] = round(float(after.get(key, 0.0)) - float(before.get(key, 0.0)), 6)
        else:
            delta[key] = int(after.get(key, 0) or 0) - int(before.get(key, 0) or 0)
    return delta


def _run_awaitable(awaitable: Any) -> Any:
    try:
        asyncio.get_running_loop()
    except RuntimeError:
        return asyncio.run(awaitable)

    result_holder: Dict[str, Any] = {}
    error_holder: Dict[str, BaseException] = {}

    def worker() -> None:
        try:
            result_holder["value"] = asyncio.run(awaitable)
        except BaseException as exc:
            error_holder["error"] = exc

    thread = threading.Thread(target=worker, daemon=True)
    thread.start()
    thread.join()

    if "error" in error_holder:
        raise error_holder["error"]
    return result_holder.get("value")


def _invoke_llm(
    *,
    prompt: str,
    model: str,
    llm_callable: Optional[Callable[..., Any]] = None,
    stage: str,
) -> str:
    if llm_callable is not None:
        try:
            result = llm_callable(prompt=prompt, model=model, stage=stage)
        except TypeError:
            result = llm_callable(prompt, model)

        if inspect.isawaitable(result):
            result = _run_awaitable(result)

        if isinstance(result, (dict, list)):
            return json.dumps(result, ensure_ascii=False)
        return str(result)

    if not _has_llm_credentials():
        raise RuntimeError(
            "Missing LLM credentials for selected provider. "
            "Use OPENAI_API_KEY for openai, or AZURE_OPENAI_API_KEY + AZURE_OPENAI_BASE_URL for azure."
        )

    response = ChatGPT_API(model=model, prompt=prompt)
    if response == "Error":
        raise RuntimeError(f"LLM call failed during stage: {stage}")
    return str(response)


def _parse_reasoning_response(raw: str) -> Dict[str, Any]:
    parsed = extract_json(raw)
    if isinstance(parsed, dict) and parsed:
        return parsed

    text = raw.strip()
    match = re.search(r"\{.*\}", text, flags=re.DOTALL)
    if match:
        try:
            obj = json.loads(match.group(0))
            if isinstance(obj, dict):
                return obj
        except Exception:
            pass

    return {"thinking": "", "node_list": []}


def build_search_tree_view(
    tree: Dict[str, Any],
    strip_fields: Sequence[str] = ("exclusive_text", "full_text"),
) -> Dict[str, Any]:
    normalized = _coerce_tree_input(tree)
    return strip_fields_fn(normalized, strip_fields)


def strip_fields_fn(data: Any, fields: Sequence[str]) -> Any:
    return strip_fields(deepcopy(data), fields)


def build_node_index(
    tree: Dict[str, Any],
    include_page_ranges: bool = True,
    exclude_fields: Optional[Sequence[str]] = None,
) -> Dict[str, Dict[str, Any]]:
    normalized = _coerce_tree_input(tree)
    rows = flatten_with_parent(normalized.get("structure", []))
    excluded = set(exclude_fields or [])

    node_index: Dict[str, Dict[str, Any]] = {}
    for row in rows:
        node_id = canonical_node_id(row.get("node_id"))
        if not node_id:
            continue

        start_index = row.get("start_index")
        end_index = row.get("end_index")
        page_span = row.get("page_span") if isinstance(row.get("page_span"), dict) else {}
        if start_index is None:
            start_index = page_span.get("start")
        if end_index is None:
            end_index = page_span.get("end")

        page_index = None
        if include_page_ranges and isinstance(start_index, int):
            if isinstance(end_index, int):
                page_index = start_index if start_index == end_index else f"{start_index}-{end_index}"
            else:
                page_index = start_index

        children = row.get("children") if isinstance(row.get("children"), list) else []
        children_ids = [
            canonical_node_id(child.get("node_id") or child.get("id"))
            for child in children
            if isinstance(child, dict)
        ]

        full_text = row.get("full_text") if isinstance(row.get("full_text"), str) else ""
        exclusive_text = row.get("exclusive_text") if isinstance(row.get("exclusive_text"), str) else ""

        row_out = {
            "node_id": node_id,
            "title": row.get("title", ""),
            "summary": row.get("summary", ""),
            "level": row.get("level"),
            "start_index": start_index,
            "end_index": end_index,
            "page_index": page_index,
            "parent_id": row.get("parent_id"),
            "children": children_ids,
            "exclusive_text": exclusive_text,
            "full_text": full_text,
        }
        if excluded:
            row_out = {key: value for key, value in row_out.items() if key not in excluded}
        node_index[node_id] = row_out
    return node_index


def reason_over_tree(
    query: str,
    tree: Dict[str, Any],
    *,
    model: Optional[str] = None,
    prompt_template: Optional[str | Callable[[str, Dict[str, Any], Dict[str, Any]], str]] = None,
    llm_callable: Optional[Callable[..., Any]] = None,
    return_prompt: bool = False,
    verbosity: Optional[str] = None,
    strip_fields: Sequence[str] = ("exclusive_text", "full_text"),
    prompt_extra: Optional[Dict[str, Any]] = None,
) -> TreeReasoningResult:
    resolved_verbosity = _normalize_verbosity(verbosity)
    logger = _configure_logger(resolved_verbosity)
    resolved_model = _resolve_model(model)
    if llm_callable is None and not resolved_model:
        raise RuntimeError(
            "Model is not configured. Pass model=... explicitly, or set "
            "OPENAI_MODEL_NAME (openai) / AZURE_DEPLOYMENT_NAME (azure)."
        )

    if resolved_verbosity in {VERBOSITY_MEDIUM, VERBOSITY_HIGH}:
        logger.info("[step] build search tree view")
    tree_view = build_search_tree_view(tree, strip_fields=strip_fields)

    if resolved_verbosity in {VERBOSITY_MEDIUM, VERBOSITY_HIGH}:
        logger.info("[step] llm tree reasoning")
    used_prompt = render_tree_search_prompt(
        query=query,
        tree_view=tree_view,
        template=prompt_template,
        extra=prompt_extra,
    )

    before = get_usage_summary()
    raw_response = _invoke_llm(
        prompt=used_prompt,
        model=resolved_model,
        llm_callable=llm_callable,
        stage="tree_search",
    )
    after = get_usage_summary()
    delta = _usage_delta(before, after)

    if resolved_verbosity == VERBOSITY_HIGH:
        logger.debug("[llm_search] model=%s", resolved_model)
        logger.debug("[llm_search] prompt_preview=%s", used_prompt[:500])
        logger.debug("[llm_search] usage_delta=%s", delta)

    parsed = _parse_reasoning_response(raw_response)
    thinking = str(parsed.get("thinking", "") or "")
    node_list_raw = parsed.get("node_list", [])
    if not isinstance(node_list_raw, list):
        node_list_raw = []

    node_list = []
    seen = set()
    for item in node_list_raw:
        normalized_id = canonical_node_id(item)
        if normalized_id and normalized_id not in seen:
            seen.add(normalized_id)
            node_list.append(normalized_id)

    return TreeReasoningResult(
        thinking=thinking,
        node_list=node_list,
        raw_response=raw_response,
        used_prompt=used_prompt if return_prompt else None,
        parsed_json=parsed,
    )


def print_reasoning_trace(
    reasoning_result: TreeReasoningResult | Dict[str, Any],
    node_index: Dict[str, Dict[str, Any]],
) -> None:
    if isinstance(reasoning_result, TreeReasoningResult):
        thinking = reasoning_result.thinking
        node_list = reasoning_result.node_list
    else:
        thinking = str(reasoning_result.get("thinking", "") or "")
        node_list = [canonical_node_id(x) for x in reasoning_result.get("node_list", []) if canonical_node_id(x)]

    print("Reasoning Process:")
    print(thinking or "<empty>")

    print("\nRetrieved Nodes:")
    for node_id in node_list:
        node = node_index.get(node_id)
        if not node:
            print(f"Node ID: {node_id}\t<missing>")
            continue
        print(
            f"Node ID: {node['node_id']}\t Page: {node.get('page_index')}\t Title: {node.get('title', '')}"
        )


def extract_selected_context(
    tree: Dict[str, Any],
    node_list: Sequence[str],
    *,
    text_mode: str = "inclusive",
    dedupe_ancestor: bool = True,
) -> ContextBundle:
    normalized = _coerce_tree_input(tree)
    node_index = build_node_index(normalized, include_page_ranges=True)

    requested_ids: List[str] = []
    seen = set()
    for value in node_list:
        node_id = canonical_node_id(value)
        if node_id and node_id not in seen:
            seen.add(node_id)
            requested_ids.append(node_id)

    missing = [node_id for node_id in requested_ids if node_id not in node_index]
    present = [node_id for node_id in requested_ids if node_id in node_index]

    dropped: List[str] = []
    selected = present
    if dedupe_ancestor:
        parent_lookup = build_parent_lookup(normalized.get("structure", []))
        selected, dropped = split_selected_with_ancestor_dedupe(present, parent_lookup)

    mode = text_mode.strip().lower()
    if mode not in {"inclusive", "exclusive"}:
        raise ValueError("text_mode must be 'inclusive' or 'exclusive'.")

    text_chunks: List[str] = []
    node_records: List[Dict[str, Any]] = []
    for node_id in selected:
        row = node_index[node_id]
        chunk = row.get("full_text", "") if mode == "inclusive" else row.get("exclusive_text", "")
        chunk = chunk if isinstance(chunk, str) else ""
        if chunk:
            text_chunks.append(chunk)

        node_records.append(
            {
                "node_id": row.get("node_id"),
                "title": row.get("title"),
                "page_index": row.get("page_index"),
                "text_length": len(chunk),
            }
        )

    return ContextBundle(
        node_list=selected,
        dropped_node_list=dropped,
        missing_node_list=missing,
        text="\n\n".join(text_chunks).strip(),
        text_mode=mode,
        node_records=node_records,
    )


def answer_from_context(
    query: str,
    context_text: str,
    *,
    model: Optional[str] = None,
    prompt_template: Optional[str | Callable[[str, str, Dict[str, Any]], str]] = None,
    llm_callable: Optional[Callable[..., Any]] = None,
    return_prompt: bool = False,
    verbosity: Optional[str] = None,
    prompt_extra: Optional[Dict[str, Any]] = None,
) -> AnswerResult:
    resolved_verbosity = _normalize_verbosity(verbosity)
    logger = _configure_logger(resolved_verbosity)
    resolved_model = _resolve_model(model)
    if llm_callable is None and not resolved_model:
        raise RuntimeError(
            "Model is not configured. Pass model=... explicitly, or set "
            "OPENAI_MODEL_NAME (openai) / AZURE_DEPLOYMENT_NAME (azure)."
        )

    used_prompt = render_answer_prompt(
        query=query,
        context_text=context_text,
        template=prompt_template,
        extra=prompt_extra,
    )

    if resolved_verbosity in {VERBOSITY_MEDIUM, VERBOSITY_HIGH}:
        logger.info("[step] generate grounded answer")

    before = get_usage_summary()
    raw_response = _invoke_llm(
        prompt=used_prompt,
        model=resolved_model,
        llm_callable=llm_callable,
        stage="answer_generation",
    )
    after = get_usage_summary()
    delta = _usage_delta(before, after)

    if resolved_verbosity == VERBOSITY_HIGH:
        logger.debug("[answer] model=%s", resolved_model)
        logger.debug("[answer] prompt_preview=%s", used_prompt[:500])
        logger.debug("[answer] usage_delta=%s", delta)

    answer = str(raw_response).strip()
    return AnswerResult(
        answer=answer,
        raw_response=str(raw_response),
        used_prompt=used_prompt if return_prompt else None,
    )


def run_reasoning_rag(
    query: str,
    tree_or_source: Any,
    *,
    model: Optional[str] = None,
    tree_prompt_template: Optional[str | Callable[[str, Dict[str, Any], Dict[str, Any]], str]] = None,
    answer_prompt_template: Optional[str | Callable[[str, str, Dict[str, Any]], str]] = None,
    llm_callable: Optional[Callable[..., Any]] = None,
    return_prompt: bool = False,
    verbosity: Optional[str] = None,
    strip_fields: Sequence[str] = ("exclusive_text", "full_text"),
    text_mode: str = "inclusive",
    dedupe_ancestor: bool = True,
    tree_prompt_extra: Optional[Dict[str, Any]] = None,
    answer_prompt_extra: Optional[Dict[str, Any]] = None,
) -> RAGResult:
    resolved_verbosity = _normalize_verbosity(verbosity)
    logger = _configure_logger(resolved_verbosity)

    tree = fetch_document_tree(tree_or_source) if not (isinstance(tree_or_source, dict) and "structure" in tree_or_source) else _coerce_tree_input(tree_or_source)

    before_total = get_usage_summary()

    reasoning = reason_over_tree(
        query=query,
        tree=tree,
        model=model,
        prompt_template=tree_prompt_template,
        llm_callable=llm_callable,
        return_prompt=return_prompt,
        verbosity=resolved_verbosity,
        strip_fields=strip_fields,
        prompt_extra=tree_prompt_extra,
    )

    node_index = build_node_index(tree, include_page_ranges=True)

    if resolved_verbosity in {VERBOSITY_MEDIUM, VERBOSITY_HIGH}:
        logger.info("[step] extract context from selected nodes")

    context = extract_selected_context(
        tree=tree,
        node_list=reasoning.node_list,
        text_mode=text_mode,
        dedupe_ancestor=dedupe_ancestor,
    )

    if resolved_verbosity == VERBOSITY_HIGH:
        logger.debug("[context] selected=%s dropped=%s missing=%s", context.node_list, context.dropped_node_list, context.missing_node_list)
        logger.debug("[context] char_count=%d", len(context.text))

    answer = answer_from_context(
        query=query,
        context_text=context.text,
        model=model,
        prompt_template=answer_prompt_template,
        llm_callable=llm_callable,
        return_prompt=return_prompt,
        verbosity=resolved_verbosity,
        prompt_extra=answer_prompt_extra,
    )

    tree_view = build_search_tree_view(tree, strip_fields=strip_fields)

    after_total = get_usage_summary()
    delta_total = _usage_delta(before_total, after_total)

    if resolved_verbosity == VERBOSITY_LOW:
        logger.info("[result] retrieved_nodes=%d context_chars=%d", len(context.node_list), len(context.text))
    elif resolved_verbosity in {VERBOSITY_MEDIUM, VERBOSITY_HIGH}:
        logger.info("[result] retrieved_nodes=%d context_chars=%d", len(context.node_list), len(context.text))
        logger.info("[result] answer_generated=yes")

    if resolved_verbosity == VERBOSITY_HIGH:
        logger.debug("[cost] cumulative_usage=%s", after_total)
        logger.debug("[cost] run_delta=%s", delta_total)

    return RAGResult(
        tree=tree,
        tree_view=tree_view,
        node_index=node_index,
        reasoning=reasoning,
        context=context,
        answer=answer,
        cost_before=before_total,
        cost_after=after_total,
        cost_delta=delta_total,
    )
