from __future__ import annotations

from app.modules.rag.contracts.enums import RagLayer
from app.modules.rag.intent_router_v2 import ConversationState, IntentRouterV2, RepoContext


def repo_context() -> RepoContext:
    return RepoContext(
        languages=["python"],
        available_domains=["CODE", "DOCS"],
        available_layers=[
            RagLayer.CODE_ENTRYPOINTS,
            RagLayer.CODE_SYMBOL_CATALOG,
            RagLayer.CODE_DEPENDENCY_GRAPH,
            RagLayer.CODE_SOURCE_CHUNKS,
            RagLayer.DOCS_MODULE_CATALOG,
            RagLayer.DOCS_FACT_INDEX,
            RagLayer.DOCS_SECTION_INDEX,
            RagLayer.DOCS_POLICY_INDEX,
        ],
    )


def run_sequence(queries: list[str], *, router: IntentRouterV2 | None = None, trace_label: str = "intent-router") -> list:
    active_router = router or IntentRouterV2()
    state = ConversationState()
    results = []
    for index, query in enumerate(queries, start=1):
        result = active_router.route(query, state, repo_context())
        print_trace(index, query, result, label=trace_label)
        results.append(result)
        state = state.advance(result)
    return results


def run_single(query: str, *, router: IntentRouterV2 | None = None, trace_label: str = "intent-router"):
    result = run_sequence([query], router=router, trace_label=trace_label)[0]
    return result


def print_trace(index: int, query: str, result, *, label: str = "intent-router") -> None:
    print(f"[{label}][turn {index}] input: {query}")
    print()
    print(f"[{label}][turn {index}] output: {result.model_dump_json(ensure_ascii=False)}")
    print("=" * 50)