from __future__ import annotations

from pathlib import Path
from threading import Lock

from app.core.api.domain.events.client_event import ClientEventRecord
from app.core.api.domain.models.agent_request import AgentRequest
from app.core.api.domain.models.agent_session import AgentSession
from app.infra.observability.trace_file_path_builder import TraceFilePathBuilder
from app.infra.observability.trace_markdown_writer import TraceMarkdownWriter


class RequestTraceLogger:
    def __init__(self, root: Path) -> None:
        self._paths = TraceFilePathBuilder(root)
        self._writer = TraceMarkdownWriter()
        self._files: dict[str, Path] = {}
        self._rag_session_files: dict[str, Path] = {}
        self._lock = Lock()

    def start_request(self, request: AgentRequest, session: AgentSession) -> None:
        path = self._resolve_path(session.active_rag_session_id)
        if not path.exists():
            self._writer.initialize(
                path,
                [
                    f"# Runtime Trace: {path.stem}",
                    "",
                    f"- active_rag_session_id: {session.active_rag_session_id or '—'}",
                ],
            )
        self._writer.append_section(
            path,
            "request",
            {
                "request_id": request.request_id,
                "session_id": session.session_id,
                "active_rag_session_id": session.active_rag_session_id,
                "process_version": request.process_version,
                "created_at": request.created_at.isoformat(),
                "message": request.message,
            },
        )
        with self._lock:
            self._files[request.request_id] = path

    def log_step(self, request_id: str, step: str, status: str, details: dict | None = None) -> None:
        self._append(request_id, f"Step {step}", {"status": status, "details": details or {}})

    def log_module(self, request_id: str, module: str, title: str, payload: dict | None = None) -> None:
        body = {"event": title}
        body.update(payload or {})
        self._append(request_id, module, body)

    def log_event(self, event: ClientEventRecord) -> None:
        return

    def complete_request(self, request: AgentRequest) -> None:
        self._append(
            request.request_id,
            "result",
            {
                "status": request.status.value,
                "answer": request.answer,
                "completed_at": request.completed_at.isoformat() if request.completed_at else None,
            },
        )

    def fail_request(self, request: AgentRequest) -> None:
        self._append(
            request.request_id,
            "result",
            {
                "status": request.status.value,
                "error": request.error.model_dump(mode="json") if request.error else None,
                "completed_at": request.completed_at.isoformat() if request.completed_at else None,
            },
        )

    def _append(self, request_id: str, title: str, payload: dict) -> None:
        with self._lock:
            path = self._files.get(request_id)
        if path is None:
            return
        self._writer.append_section(path, title, payload)

    def _resolve_path(self, rag_session_id: str | None) -> Path:
        if not rag_session_id:
            return self._paths.build_for_rag_session(None)
        with self._lock:
            path = self._rag_session_files.get(rag_session_id)
            if path is None:
                path = self._paths.build_for_rag_session(rag_session_id)
                self._rag_session_files[rag_session_id] = path
            return path
