"""Tatasteel ASGI server — HTTP query endpoint + chat UI.

This module owns the FastAPI/Starlette app lifecycle and the REST handler.
The agent definition lives in ``backend.chanakya.tatasteel.agent``.

Entrypoint::

    uvicorn backend.chanakya.tatasteel.server:app --host 0.0.0.0 --port 8005
"""
import asyncio
import os
from pathlib import Path

import asyncpg
import logfire
from loguru import logger as _log
from starlette.requests import Request
from starlette.responses import JSONResponse

from backend.chanakya.a2a_context import make_metadata_aware_app
from backend.chanakya.config import AGENT_REGISTRY
from backend.chanakya.conversation import deserialize_history, load_conversation_history
from backend.chanakya.nl2sql_core import NL2SQLDeps
from backend.chanakya.schemas import AgentQueryRequest, AgentQueryResponse, UIRenderRequest
from backend.chanakya.tatasteel.agent import tatasteel_agent
from pylogue.shell import app_factory as create_core_app

_cfg = AGENT_REGISTRY["tatasteel"]
_PROJECT_ROOT = Path(__file__).parents[3]
_DEFAULT_SCHEMA_PATH = Path(__file__).parent / "schema.sql"

# ---------------------------------------------------------------------------
# ASGI app
# ---------------------------------------------------------------------------
app = make_metadata_aware_app(
    tatasteel_agent,
    name=_cfg.name,
    description=_cfg.description,
)

_db_pool: asyncpg.Pool | None = None


async def _get_pool() -> asyncpg.Pool:
    global _db_pool
    if _db_pool is None:
        _db_pool = await asyncpg.create_pool(
            host=os.environ.get("TATASTEEL_DB_HOST", "localhost"),
            port=int(os.environ.get("TATASTEEL_DB_PORT", "5432")),
            database=os.environ["TATASTEEL_DB_NAME"],
            user=os.environ["TATASTEEL_DB_USER"],
            password=os.environ["TATASTEEL_DB_PASSWORD"],
            min_size=2,
            max_size=10,
        )
    return _db_pool


# ---------------------------------------------------------------------------
# REST endpoint — structured JSON
# ---------------------------------------------------------------------------

async def _query_handler(request: Request) -> JSONResponse:
    """POST /query — runs the SQL query and returns raw rows + explanation.

    The response is shaped as a ``UIRenderRequest`` so the caller (Chanakya or
    the frontend) can forward it directly to the UI agent's ``POST /query``.
    """
    try:
        body = await request.json()
        req = AgentQueryRequest(**body)
    except Exception as exc:
        return JSONResponse(
            AgentQueryResponse.failure(message=f"Invalid request: {exc}", code="PIPELINE_ERROR").model_dump(),
            status_code=400,
        )

    if req.history is not None:
        message_history = deserialize_history(req.history)
        _log.debug("[TATASTEEL] Using {} wire-format history messages (skipped DB)", len(message_history))
    else:
        message_history = await load_conversation_history(req.conversation_id)
        _log.debug("[TATASTEEL] Loaded {} messages for conversation {}", len(message_history), req.conversation_id)

    pool = await _get_pool()
    deps = NL2SQLDeps(user_email=req.user_email or "", progress_queue=asyncio.Queue(), pool=pool)

    try:
        with logfire.span("{q}", q=req.question[:120]):
            result = await tatasteel_agent.agent.run(req.question, deps=deps, message_history=message_history)
    except Exception as exc:
        _log.error("tatasteel query error: {}", exc)
        return JSONResponse(
            AgentQueryResponse.failure(message=str(exc), code="PIPELINE_ERROR").model_dump()
        )

    output = result.output
    progress_steps: list[dict] = []
    q = deps.progress_queue
    while q is not None and not q.empty():
        progress_steps.append(q.get_nowait())

    rows = deps.query_result or []
    sql = deps.generated_sql or ""
    rows_sample = rows[:3] if rows else []

    return JSONResponse(
        {
            **UIRenderRequest(
                question=req.question,
                explanation=output.explanation,
                rows=rows,
                user_email=req.user_email,
                conversation_id=req.conversation_id,
                progress_steps=progress_steps,
                generated_sql=sql,
                sql_explanation=deps.sql_explanation,
                what_to_act_on=output.what_to_act_on or None,
                urgency=output.urgency or None,
                cost_of_inaction=output.cost_of_inaction or None,
            ).model_dump(),
            "generated_sql": sql,
            "rows_sample": rows_sample,
            "query_log": deps.query_log,
        }
    )


app.add_route("/query", _query_handler, methods=["POST"])

# ---------------------------------------------------------------------------
# Startup banner + chat UI
# ---------------------------------------------------------------------------

@app.on_event("startup")
async def _startup() -> None:
    port = os.environ.get("TATASTEEL_AGENT_PORT", "8005")
    db = os.environ.get("TATASTEEL_DB_NAME", "(not set)")
    host = os.environ.get("TATASTEEL_DB_HOST", "localhost")
    _log.info("\n  Tatasteel Agent ready:")
    _log.info("    A2A   → http://0.0.0.0:{}/", port)
    _log.info("    Chat  → http://0.0.0.0:{}/chat", port)
    _log.info("    Query → http://0.0.0.0:{}/query  (returns UIRenderRequest for UI agent)", port)
    _log.info("    DB    → postgres://{}:{}/{}", host, os.environ.get("TATASTEEL_DB_PORT", "5432"), db)
    _log.info("    Schema→ {}", _DEFAULT_SCHEMA_PATH)
    await _get_pool()
    _log.info("    DB pool → ready (min=2, max=10)")


@app.on_event("shutdown")
async def _shutdown() -> None:
    global _db_pool
    if _db_pool is not None:
        await _db_pool.close()
        _db_pool = None


_TATASTEEL_DB_PATH = os.environ.get(
    "PYLOGUE_DB_PATH",
    str(_PROJECT_ROOT / "data" / "conversation-histories" / "tatasteel" / "chat_history.db"),
)
Path(_TATASTEEL_DB_PATH).parent.mkdir(parents=True, exist_ok=True)

app.mount(
    "/chat",
    create_core_app(
        responder_factory=lambda: tatasteel_agent,
        hero_title="Tatasteel Agent",
        hero_subtitle="Query the Tatasteel construction project database with natural language.",
        db_path=_TATASTEEL_DB_PATH,
    ),
    name="tatasteel-chat",
)


if __name__ == "__main__":
    import uvicorn

    uvicorn.run("backend.chanakya.tatasteel.server:app", host="0.0.0.0", port=_cfg.port, reload=True)