app.py

from __future__ import annotations

import re
from pathlib import Path

import chainlit as cl
from loguru import logger

from config import CLAIM_PACKETS_ROOT
from graphs import claims_graph, scheduling_graph

CLAIM_DEMO_SCENARIOS = [
    {
        "claim_id": "CLM-0001",
        "scenario": "Clean professional claim",
        "expected": "Clean pass; routed to clean_pass_auto_normalization",
        "details": "MBR100239 | Bayview Orthopedics | CPT 99213 | ICD-10 M25.561",
    },
    {
        "claim_id": "CLM-0002",
        "scenario": "MRI likely needs prior authorization",
        "expected": "Prior auth risk; missing priorAuthNumber; routed to prior_auth_exception_review",
        "details": "MBR100548 | Golden State Imaging | CPT 73721 | ICD-10 M25.561",
    },
    {
        "claim_id": "CLM-0003",
        "scenario": "Coverage inactive on DOS",
        "expected": "Eligibility gap; routed to claims_ops_exception_review",
        "details": "MBR100791 | East Valley Cardiology | CPT 99214 | ICD-10 E11.9",
    },
    {
        "claim_id": "CLM-0004",
        "scenario": "Diagnosis/procedure mismatch",
        "expected": "Coding mismatch risk; routed to claims_ops_exception_review",
        "details": "MBR101044 | Sunrise PT Group | CPT 93000 | ICD-10 S83.241A",
    },
    {
        "claim_id": "CLM-0005",
        "scenario": "Possible duplicate claim",
        "expected": "Duplicate risk against CLM-0006; routed to claims_ops_duplicate_review",
        "details": "MBR100239 | Bayview Orthopedics | CPT 97110 | ICD-10 M54.50",
    },
    {
        "claim_id": "CLM-0006",
        "scenario": "Possible duplicate claim partner",
        "expected": "Duplicate risk against CLM-0005; routed to claims_ops_duplicate_review",
        "details": "MBR100548 | Golden State Imaging | CPT 97110 | ICD-10 M54.50",
    },
    {
        "claim_id": "CLM-0007",
        "scenario": "Attachment control number mismatch",
        "expected": "Attachment/documentation mismatch; routed to claims_ops_exception_review",
        "details": "MBR100791 | East Valley Cardiology | CPT 27447 | ICD-10 M17.11",
    },
    {
        "claim_id": "CLM-0008",
        "scenario": "Out-of-network specialist referral",
        "expected": "OON/referral risk; routed to claims_ops_exception_review",
        "details": "MBR101044 | Sunrise PT Group | CPT 99204 | ICD-10 I48.91",
    },
]


SCHEDULING_DEMO_SCENARIOS = [
    {
        "case_id": "C9001",
        "referral_id": "R5001",
        "member_id": "M1001",
        "specialty": "Cardiology",
        "request": "Schedule member M1001 with an in-network cardiology specialist near San Ramon next week. This is urgent and SMS confirmation is preferred.",
        "desired_outcome": "schedule earliest in-network cardiology visit within 7 days",
        "constraints": "HMO referral active; urgent; San Ramon preferred; SMS confirmation",
        "expected": "Schedule earliest in-network cardiology visit within 7 days using active HMO referral R5001; urgent; San Ramon preferred; SMS confirmation.",
    },
    {
        "case_id": "C9002",
        "referral_id": "R5004",
        "member_id": "M1004",
        "specialty": "Cardiology",
        "request": "Schedule member M1004 for a post-discharge in-network cardiology visit before the authorization expires. The location must be wheelchair accessible and within 10 miles if possible, with phone confirmation.",
        "desired_outcome": "schedule post-discharge cardiology before authorization expires",
        "constraints": "wheelchair accessible; max travel 10 miles; phone confirmation",
        "expected": "Schedule post-discharge cardiology for M1004 before authorization A7003 expires; wheelchair-accessible location within 10 miles; phone confirmation.",
    },
    {
        "case_id": "C9003",
        "referral_id": "R5005",
        "member_id": "M1005",
        "specialty": "OB-GYN",
        "request": "Schedule member M1005 with an in-network OB-GYN for high-risk pregnancy care, but only if prior authorization is approved and required documentation is complete.",
        "desired_outcome": "do not schedule until PA approved or missing documentation collected",
        "constraints": "EPO in-network only; high-risk OB-GYN PA pending",
        "expected": "Do not schedule until PA is approved and documentation is complete; EPO in-network only; high-risk OB-GYN referral R5005 / auth A7004 still pending.",
    },
    {
        "case_id": "C9004",
        "referral_id": "R5006",
        "member_id": "M1006",
        "specialty": "Neurology",
        "request": "Schedule member M1006 with a neurology specialist. If the closest neurologist is not accepting new patients, route this as an exception.",
        "desired_outcome": "route exception because closest neurologist is not accepting new patients",
        "constraints": "PPO plan; neurology requested; provider new-patient status false",
        "expected": "Route as exception when the closest neurologist is not accepting new patients; PPO plan with neurology referral R5006.",
    },
    {
        "case_id": "C9005",
        "referral_id": "R5008",
        "member_id": "M1008",
        "specialty": "Behavioral Health",
        "request": "Schedule member M1008 for a telehealth behavioral health or psychiatry evaluation, preferably late afternoon.",
        "desired_outcome": "schedule telehealth psychiatry/behavioral health evaluation",
        "constraints": "late afternoon; telehealth acceptable; Student PPO",
        "expected": "Schedule telehealth behavioral health/psychiatry evaluation for M1008; prefer late afternoon; Student PPO with referral R5008.",
    },
]


def find_scheduling_scenario(result: dict) -> dict | None:
    """Match a graph result back to a demo scenario by member id or request text."""
    member_id = result.get("member_id") or result.get("extracted_request", {}).get(
        "member_id"
    )
    if member_id:
        for item in SCHEDULING_DEMO_SCENARIOS:
            if item.get("member_id") == member_id:
                return item

    request_text = (
        result.get("request_text")
        or result.get("extracted_request", {}).get("raw_text")
        or ""
    ).strip()
    if request_text:
        for item in SCHEDULING_DEMO_SCENARIOS:
            if item.get("request", "").strip() == request_text:
                return item
    return None


def render_claim_result(result: dict) -> str:
    """Render scenario-specific claims output for Chainlit."""
    claim = result.get("canonical_claim", {})
    validation = result.get("validation_results", {})
    final = result.get("final_summary", {})
    risk = result.get("denial_risk") or final.get("risk", {})
    route = result.get("route") or final.get("route", {})
    llm_decision = result.get("llm_decision") or final.get("llm_decision", {})
    rag_decision = result.get("rag_decision") or final.get("rag_decision", {})

    risks = risk.get("risks", []) if isinstance(risk, dict) else []
    risk_rows = (
        "\n".join(f"- `{r.get('risk')}` — {r.get('severity')}" for r in risks)
        or "- None detected"
    )

    cpts = ", ".join(claim.get("service", {}).get("cpt_codes", []) or [])
    icds = ", ".join(claim.get("service", {}).get("icd10_codes", []) or [])
    dates = ", ".join(claim.get("service", {}).get("dates", []) or [])

    return f"""
## Claims workflow result: `{claim.get('claim_id', result.get('case_id', 'unknown'))}`

| Field | Value |
|---|---|
| Member | `{claim.get('member', {}).get('member_id')}` |
| Provider | {claim.get('provider', {}).get('name')} |
| Network status | `{claim.get('provider', {}).get('network_status')}` |
| CPT | `{cpts}` |
| ICD-10 | `{icds}` |
| Date of service | `{dates}` |
| Eligibility | `{validation.get('eligible')}` |
| Provider valid | `{validation.get('provider_valid')}` |
| Duplicate risk | `{validation.get('duplicate_risk')}` |
| Risk level | `{risk.get('risk_level', llm_decision.get('risk_level', 'unknown'))}` |
| Route | `{route.get('route', llm_decision.get('recommended_route', 'unknown'))}` |

### Risks detected
{risk_rows}

### Retrieval / decision details
- Policy RAG results: `{final.get('policy_context_count', len(result.get('policy_results', {}).get('results', [])))}`
- Similar exception results: `{final.get('similar_exception_count', len(result.get('exception_results', {}).get('results', [])))}`
- RAG decision: {rag_decision.get('reason', 'not available')}

### Audit explanation
{llm_decision.get('explanation', 'Deterministic tools completed the claim workflow.')}
""".strip()


def _format_appointment_option(option: dict, index: int) -> str:
    """Format a ranked option whether it came from state or a summarized payload."""
    slot = option.get("slot")
    if isinstance(slot, dict):
        start = slot.get("start")
        location = slot.get("location")
        city = slot.get("city")
        provider = slot.get("provider")
    else:
        start = option.get("start_datetime") or option.get("start")
        location = option.get("location_name") or option.get("location")
        if isinstance(location, dict):
            location = location.get("location_name")
        city = option.get("city")
        provider = option.get("provider_name") or option.get("provider")

    rank = option.get("rank", index + 1)
    if not any([start, location, city, provider]):
        return f"- Option {rank}: appointment details unavailable"
    return f"- Option {rank}: {start} at {location} ({city}) with {provider}"


def _scheduling_appointment_options(result: dict) -> list:
    """Prefer graph state options; final_summary may contain LLM placeholders."""
    state_options = result.get("appointment_options") or []
    final_options = result.get("final_summary", {}).get("appointment_options") or []

    def has_details(option: dict) -> bool:
        slot = option.get("slot")
        if isinstance(slot, dict) and any(slot.values()):
            return True
        return bool(
            option.get("start_datetime")
            or option.get("start")
            or option.get("location_name")
            or option.get("provider_name")
        )

    if state_options and any(has_details(option) for option in state_options):
        return state_options
    if final_options and any(has_details(option) for option in final_options):
        return final_options
    return state_options or final_options


def render_scheduling_result(result: dict) -> str:
    """Render scenario-specific scheduling output for Chainlit."""
    final = result.get("final_summary", {})
    extracted = result.get("extracted_request", {})
    readiness = result.get("schedule_readiness", {})
    scenario = find_scheduling_scenario(result)
    member_id = result.get("member_id") or extracted.get("member_id")
    options = _scheduling_appointment_options(result)
    option_rows = (
        "\n".join(
            _format_appointment_option(option, i) for i, option in enumerate(options)
        )
        or "- No appointment options returned"
    )
    issues = (
        "\n".join(f"- {issue}" for issue in readiness.get("issues", [])) or "- None"
    )
    return f"""
## Scheduling workflow result

| Field | Value |
|---|---|
| Case | `{scenario['case_id'] if scenario else 'n/a'}` |
| Member | `{member_id}` |
| Specialty | `{extracted.get('specialty') or (scenario or {}).get('specialty', '')}` |
| Preferred city | `{extracted.get('city') or 'any'}` |
| Ready to schedule | `{readiness.get('ready_to_schedule')}` |
| Recommended action | `{final.get('recommended_action')}` |

### Readiness issues
{issues}

### Appointment options
{option_rows}

### Summary
{final.get('member_facing_summary', 'Deterministic scheduling workflow completed.')}
""".strip()


def render_claim_demo_scenarios() -> str:
    rows = [
        "| Claim to enter | Scenario | Expected outcome |",
        "|---|---|---|",
    ]
    for item in CLAIM_DEMO_SCENARIOS:
        rows.append(
            f"| `{item['claim_id']}` | {item['scenario']} | {item['expected']} |"
        )
    return "\n".join(rows)


def render_scheduling_demo_scenarios() -> str:
    rows = [
        "| Case | Member | Request | Expected outcome |",
        "|---|---|---|---|",
    ]
    for item in SCHEDULING_DEMO_SCENARIOS:
        rows.append(
            f"| {item['case_id']} | {item['member_id']} | `{item['request']}` | {item['expected']} |"
        )
    return "\n".join(rows)


def resolve_claim_packet_path(user_input: str) -> Path:
    """Resolve a claim number like CLM-0004 to the local packet folder path."""
    value = user_input.strip()

    # If the user still pastes a full path, keep supporting it.
    direct_path = Path(value)
    if direct_path.exists():
        return direct_path

    normalized = value.upper().replace("_", "-")
    match = re.search(r"CLM-?0*(\d+)", normalized)
    if match:
        claim_number = int(match.group(1))
        canonical = f"CLM-{claim_number:04d}"
        candidates = [
            CLAIM_PACKETS_ROOT / canonical,
            CLAIM_PACKETS_ROOT / canonical.lower(),
            CLAIM_PACKETS_ROOT / canonical.replace("-", "_"),
            CLAIM_PACKETS_ROOT / canonical.replace("-", "_").lower(),
            CLAIM_PACKETS_ROOT / f"clm_{claim_number:03d}",
            CLAIM_PACKETS_ROOT / f"clm-{claim_number:04d}",
        ]
        for candidate in candidates:
            if candidate.exists():
                return candidate
        return CLAIM_PACKETS_ROOT / canonical

    # Fallback: append raw input to the packet root.
    return CLAIM_PACKETS_ROOT / value


@cl.on_chat_start
async def on_chat_start():
    cl.user_session.set("workflow", None)
    actions = [
        cl.Action(
            name="claims",
            payload={},
            label="Prototype 1: Claims Ingestion & Normalization",
        ),
        cl.Action(
            name="scheduling",
            payload={},
            label="Prototype 2: Scheduling & Admin Assistant",
        ),
    ]
    await cl.Message(
        content="Choose a prototype to run.",
        actions=actions,
    ).send()


@cl.action_callback("claims")
async def on_claims(action: cl.Action):
    cl.user_session.set("workflow", "claims")

    content = f"""
Claims workflow selected.

Enter only the claim number to run a scenario — for example `CLM-0004`.

{render_claim_demo_scenarios()}
"""

    await cl.Message(content=content).send()


@cl.action_callback("scheduling")
async def on_scheduling(action: cl.Action):
    cl.user_session.set("workflow", "scheduling")

    content = f"""
Scheduling workflow selected.

Copy/paste one of the sample scheduling requests below.

{render_scheduling_demo_scenarios()}
"""

    await cl.Message(content=content).send()


@cl.on_message
async def on_message(message: cl.Message):
    workflow = cl.user_session.get("workflow")
    if not workflow:
        await on_chat_start()
        return

    if workflow == "claims":
        claim_input = message.content.strip()
        packet_path = resolve_claim_packet_path(claim_input)
        logger.info(
            f"Running claims graph for claim_input={claim_input}, packet_path={packet_path}"
        )

        if not packet_path.exists():
            await cl.Message(
                content=(
                    f"Claim packet not found for `{claim_input}`.\n\n"
                    f"I looked under: `{packet_path}`\n\n"
                    "Enter a claim number such as `CLM-0004`, or verify the packet folder exists under "
                    f"`{CLAIM_PACKETS_ROOT}`."
                )
            ).send()
            return

        result = claims_graph.invoke(
            {"selected_packet_path": str(packet_path), "messages": []}
        )
        await cl.Message(content="Claims workflow complete.").send()
        await cl.Message(content=render_claim_result(result)).send()
        return

    if workflow == "scheduling":
        logger.info("Running scheduling graph")
        result = scheduling_graph.invoke(
            {"request_text": message.content, "messages": []}
        )
        await cl.Message(content="Scheduling workflow complete.").send()
        await cl.Message(content=render_scheduling_result(result)).send()
        return