from __future__ import annotations

from inspect import isclass
from typing import Any, Literal

from aiel_sdk import AielClient, http
from aiel_sdk.adapters.langgraph import AIELLangGraphMemory
from aiel_sdk.integrations.openai import ChatOpenAI
from aiel_sdk.integrations.postgres.remote import RemotePostgres
from aiel_sdk.integrations.postgres.types import PostgresInsertRequest, PostgresQueryRequest
from aiel_sdk.memory import RemoteSaver

from aiel_core.langgraph.graph import END, START, StateGraph
from aiel_core.aiel_cli.context.user_context import UserContext
try:
    from aiel_runtime.protocol import Context as AielContext
except Exception:  # pragma: no cover
    AielContext = object  # type: ignore[assignment]

from .types import SupportState


@http.post("/support/case")
def http_support_case(ctx, body: dict) -> dict:
    ctx.log("http.support.case.post", body=body)
    return body


@http.get("/support/case")
def http_support_case_get(ctx, req) -> dict:
    safe_req = {"keys": list(req.keys())} if isinstance(req, dict) else {"type": str(type(req))}
    return {"status": "OK", "request_id": getattr(ctx, "request_id", None), "req": safe_req}


ctx = UserContext().load()


class OpenAIIntegration:
    def __init__(self, connection_id: str, model: str = "gpt-4o-mini") -> None:
        self.workspace_id = ctx.workspace_id
        self.project_id = ctx.project_id
        self._client = AielClient(base_url=ctx.base_url, api_key=ctx.token)
        self.openai = ChatOpenAI(
            client=self._client,
            workspace_id=self.workspace_id,
            project_id=self.project_id,
            connection_id=connection_id,
            model=model,
            temperature=0.2,
        )

    def responses_create(self, *, system: str, prompt: str) -> str:
        raw = self.openai.invoke(
            [
                {"role": "system", "content": system},
                {"role": "user", "content": prompt},
            ]
        )
        choices = raw.get("choices")
        if isinstance(choices, list) and choices and isinstance(choices[0], dict):
            msg = choices[0].get("message")
            if isinstance(msg, dict) and isinstance(msg.get("content"), str):
                return msg["content"].strip()
        output_text = raw.get("output_text")
        if isinstance(output_text, str):
            return output_text.strip()
        return str(raw)


class PostgresIntegration:
    def __init__(self, connection_id: str) -> None:
        self.workspace_id = ctx.workspace_id
        self.project_id = ctx.project_id
        self._client = AielClient(base_url=ctx.base_url, api_key=ctx.token)
        self.postgres = RemotePostgres(
            client=self._client,
            workspace_id=self.workspace_id,
            project_id=self.project_id,
            connection_id=connection_id,
        )

    def query(self, sql: str, params: dict[str, Any] | None = None, max_rows: int = 50) -> list[dict[str, Any]]:
        out = self.postgres.query(PostgresQueryRequest(sql=sql, params=params or {}, max_rows=max_rows))
        return out.rows

    def insert(self, table: str, values: dict[str, Any], returning: list[str] | None = None) -> dict[str, Any]:
        out = self.postgres.insert(PostgresInsertRequest(table=table, values=values, returning=returning or []))
        return {"inserted": out.inserted, "returning_rows": out.returning_rows, "raw": out.raw}


def _build_memory() -> AIELLangGraphMemory:
    client = AielClient(base_url=ctx.base_url, api_key=ctx.token)
    saver = RemoteSaver(client, workspace_id=ctx.workspace_id, project_id=ctx.project_id)
    return AIELLangGraphMemory(saver)


def route_after_triage(state: SupportState) -> Literal["research_agent", "response_agent"]:
    return "research_agent" if state.get("route") == "research" else "response_agent"


async def triage_agent(state: SupportState, context: Any) -> SupportState:
    openai_resource_id = str(state.get("openai_resource_id") or "").strip()
    if not openai_resource_id:
        raise RuntimeError("openai_resource_id is required")
    openai = OpenAIIntegration(openai_resource_id)
    request_text = state.get("request_text") or ""
    summary = openai.responses_create(
        system=(
            "You are a support triage agent. Output compact JSON with keys: "
            "summary (string), label (billing|general), route (research|respond)."
        ),
        prompt=request_text,
    )
    lowered = request_text.lower()
    is_billing = any(k in lowered for k in ("invoice", "charged", "refund", "billing", "payment"))
    return {
        "triage_summary": summary,
        "triage_label": "billing" if is_billing else "general",
        "route": "research" if is_billing else "respond",
    }


async def research_agent(state: SupportState, context: Any) -> SupportState:
    _ = context
    postgres_resource_id = str(state.get("postgres_resource_id") or "").strip()
    if not postgres_resource_id:
        raise RuntimeError("postgres_resource_id is required")
    pg = PostgresIntegration(postgres_resource_id)
    account_id = state.get("account_id") or "acct_001"
    customer_rows = pg.query(
        sql=(
            "SELECT account_id, plan_name, mrr, status, renewal_date "
            "FROM customer_accounts WHERE account_id = %(account_id)s LIMIT 1"
        ),
        params={"account_id": account_id},
        max_rows=1,
    )
    invoice_rows = pg.query(
        sql=(
            "SELECT invoice_id, due_date, amount, status "
            "FROM invoices WHERE account_id = %(account_id)s "
            "ORDER BY due_date DESC LIMIT 5"
        ),
        params={"account_id": account_id},
        max_rows=5,
    )
    return {
        "account_id": account_id,
        "customer_360": {
            "customer": customer_rows[0] if customer_rows else {},
            "invoices": invoice_rows,
        },
    }


async def risk_agent(state: SupportState, context: Any) -> SupportState:
    _ = context
    postgres_resource_id = str(state.get("postgres_resource_id") or "").strip()
    if not postgres_resource_id:
        raise RuntimeError("postgres_resource_id is required")
    pg = PostgresIntegration(postgres_resource_id)
    account_id = state["account_id"]
    disputes = pg.query(
        sql=(
            "SELECT dispute_id, amount, status "
            "FROM payment_disputes WHERE account_id = %(account_id)s "
            "ORDER BY created_at DESC LIMIT 10"
        ),
        params={"account_id": account_id},
        max_rows=10,
    )
    risk_score = min(100, len(disputes) * 20)
    risk_level = "high" if risk_score >= 70 else ("medium" if risk_score >= 35 else "low")
    return {"risk_report": {"risk_score": risk_score, "risk_level": risk_level, "recent_disputes": disputes}}


async def decision_agent(state: SupportState, context: Any) -> SupportState:
    _ = context
    openai_resource_id = str(state.get("openai_resource_id") or "").strip()
    if not openai_resource_id:
        raise RuntimeError("openai_resource_id is required")
    openai = OpenAIIntegration(openai_resource_id)
    risk = state.get("risk_report") or {}
    invoices = (state.get("customer_360") or {}).get("invoices") or []
    risk_level = str(risk.get("risk_level") or "low")
    if risk_level == "high":
        action = "escalate_to_finance"
    else:
        has_failed = any(str(inv.get("status") or "").lower() in {"overdue", "failed"} for inv in invoices)
        action = "offer_payment_plan" if has_failed else "issue_credit_note"

    reasoning = openai.responses_create(
        system="You are a support decision agent. Explain the action in 2 concise sentences.",
        prompt=(
            f"Triage:\n{state.get('triage_summary')}\n\n"
            f"Risk:\n{risk}\n\n"
            f"Chosen action: {action}"
        ),
    )
    return {"decision": {"action": action, "risk_level": risk_level, "reasoning": reasoning}}


async def execution_agent(state: SupportState, context: Any, *, store: Any) -> SupportState:
    _ = context
    postgres_resource_id = str(state.get("postgres_resource_id") or "").strip()
    if not postgres_resource_id:
        raise RuntimeError("postgres_resource_id is required")
    pg = PostgresIntegration(postgres_resource_id)
    decision = state.get("decision") or {}
    inserted = pg.insert(
        table="support_case_actions",
        values={
            "case_id": state.get("case_id"),
            "account_id": state.get("account_id"),
            "user_id": state.get("user_id"),
            "action": decision.get("action"),
            "risk_level": decision.get("risk_level"),
            "reasoning": decision.get("reasoning"),
        },
        returning=["id", "created_at"],
    )
    # persist long-term memory in store
    user_id = str(state.get("user_id") or "")
    if user_id:
        from langgraph.store.base.batch import PutOp

        store.batch(
            [
                PutOp(
                    namespace=("support", "users", user_id),
                    key="latest_resolution_preference",
                    value={
                        "last_action": decision.get("action"),
                        "risk_level": decision.get("risk_level"),
                    },
                )
            ]
        )
    return {"execution_result": inserted}


async def response_agent(state: SupportState, context: Any, *, store: Any) -> SupportState:
    _ = context
    openai_resource_id = str(state.get("openai_resource_id") or "").strip()
    if not openai_resource_id:
        raise RuntimeError("openai_resource_id is required")
    openai = OpenAIIntegration(openai_resource_id)
    decision = state.get("decision") or {}
    execution = state.get("execution_result") or {}
    user_id = str(state.get("user_id") or "")
    preference = None
    if user_id:
        from langgraph.store.base.batch import GetOp
            #aecio esteve aqui
        preference = store.batch([GetOp(namespace=("support", "users", user_id), key="latest_resolution_preference")])[0]
    final_text = openai.responses_create(
        system="You are a customer support response agent. Return a short professional response.",
        prompt=(
            f"Customer request: {state.get('request_text')}\n"
            f"Triage summary: {state.get('triage_summary')}\n"
            f"Decision: {decision}\n"
            f"Execution result: {execution}\n"
            f"User long-memory preference: {preference.value if preference else {}}"
        ),
    )
    return {"final_response": final_text}


async def aiel_handler(event: SupportState, context):
    print(isclass(AielContext), (AielContext, type(AielContext)))
    memory = _build_memory()
    checkpointer = memory.checkpointer
    store = memory.store

    graph = StateGraph(SupportState, context_schema=AielContext)
    graph.add_node("triage_agent", triage_agent)
    graph.add_node("research_agent", research_agent)
    graph.add_node("risk_agent", risk_agent)
    graph.add_node("decision_agent", decision_agent)
    graph.add_node("execution_agent", execution_agent)
    graph.add_node("response_agent", response_agent)

    graph.add_edge(START, "triage_agent")
    graph.add_conditional_edges("triage_agent", route_after_triage)
    graph.add_edge("research_agent", "risk_agent")
    graph.add_edge("risk_agent", "decision_agent")
    graph.add_edge("decision_agent", "execution_agent")
    graph.add_edge("execution_agent", "response_agent")
    graph.add_edge("response_agent", END)

    app = graph.compile(checkpointer=checkpointer, store=store)
    thread_id = event.get("case_id") or "support-default-thread"
    result = await app.ainvoke(
        event,
        context=context,
        config={"configurable": {"thread_id": thread_id, "checkpoint_ns": "support-runtime"}},
    )
    return {"result": result, "context": context}
