"""
AJCC Anatomic Staging API (Simplified)
--------------------------------------
Single-file FastAPI app that exposes TNM -> Stage mapping for:

1) Breast Cancer – AJCC 8 (Anatomic)
2) Lung Cancer (NSCLC) – AJCC 8
3) Colon / Rectal Cancer – AJCC 8
4) Stomach (Gastric Cancer) – AJCC 8
5) Pancreatic Cancer – AJCC 8
6) Liver Cancer (Hepatocellular Carcinoma) – AJCC 8
7) Melanoma – AJCC 8
8) Prostate Cancer – (Anatomic logic only)
9) Kidney (Renal Cell Carcinoma)
10) Bladder Cancer – AJCC 8

Plus a general fallback AJCC rule for any other / unknown cancer types.

This is a simplified logic based on the mapping you provided,
NOT a full official AJCC implementation.
"""

from typing import Optional, Any, Dict, List

from fastapi import FastAPI
from pydantic import BaseModel


# ---------------------------------------------------------------------
#  Core TNM staging logic (previously in ajcc_staging.py)
# ---------------------------------------------------------------------

def normalize_tnm(value: Optional[str]) -> Optional[str]:
    if value is None:
        return None
    value = value.strip()
    if not value:
        return None
    return value.upper()


def stage_cancer(cancer_type: str, T: str, N: str, M: str) -> str:
    """
    Main entry point.

    cancer_type: string key identifying the cancer
    T, N, M: TNM strings like 'T1', 'T2B', 'N1', 'N1MI', 'M1A', etc.

    Returns: stage string like 'Stage IIA', 'Stage IIIC', 'Stage IVB', etc.
    """
    cancer_type = cancer_type.strip().lower()
    T = normalize_tnm(T)
    N = normalize_tnm(N)
    M = normalize_tnm(M)

    if cancer_type == "breast":
        return stage_breast(T, N, M)

    elif cancer_type in ("lung", "nsclc", "non-small cell lung", "non small cell lung"):
        return stage_nsclc(T, N, M)

    elif cancer_type in ("colon", "rectal", "colorectal"):
        return stage_colorectal(T, N, M)

    elif cancer_type in ("stomach", "gastric"):
        return stage_stomach(T, N, M)

    elif cancer_type in ("pancreas", "pancreatic"):
        return stage_pancreas(T, N, M)

    elif cancer_type in ("liver", "hcc", "hepatocellular"):
        return stage_liver(T, N, M)

    elif cancer_type == "melanoma":
        return stage_melanoma(T, N, M)

    elif cancer_type == "prostate":
        return stage_prostate(T, N, M)

    elif cancer_type in ("kidney", "renal"):
        return stage_kidney(T, N, M)

    elif cancer_type == "bladder":
        return stage_bladder(T, N, M)

    else:
        return stage_general(T, N, M)

    

# ---------------------------------------------------------------------
# General fallback AJCC rule (for unknown cancer types)
# ---------------------------------------------------------------------

def stage_general(T: str, N: str, M: str) -> str:
    """
    General AJCC mapping:

    IF T = Tis AND N = N0 AND M = M0 -> Stage 0
    IF (T = T1 OR T = T2) AND N = N0 AND M = M0 -> Stage I
    IF (T = T3 AND N = N0 AND M = M0)
       OR ((T = T1 OR T = T2) AND N = N1 AND M = M0) -> Stage II
    IF ((T = T3 OR T = T4) AND (N = N1 OR N = N2 OR N = N3) AND M = M0)
       OR ((T = T1 OR T = T2 OR T = T3 OR T = T4)
           AND (N = N2 OR N = N3) AND M = M0) -> Stage III
    IF M = M1 -> Stage IV
    """

    # Stage IV
    if M == "M1":
        return "Stage IV"

    # Stage 0
    if T == "TIS" and N == "N0" and M == "M0":
        return "Stage 0"

    # Stage I
    if T in ("T1", "T2") and N == "N0" and M == "M0":
        return "Stage I"

    # Stage II
    if (T == "T3" and N == "N0" and M == "M0") or (
        T in ("T1", "T2") and N == "N1" and M == "M0"
    ):
        return "Stage II"

    # Stage III
    if (
        (T in ("T3", "T4") and N in ("N1", "N2", "N3") and M == "M0")
        or (T in ("T1", "T2", "T3", "T4") and N in ("N2", "N3") and M == "M0")
    ):
        return "Stage III"

    return "Stage not mapped (general)"


# ---------------------------------------------------------------------
# 1) Breast Cancer – AJCC 8 (Anatomic)
# ---------------------------------------------------------------------

def stage_breast(T: str, N: str, M: str) -> str:
    # Stage IV first
    if M == "M1":
        return "Stage IV"

    # Stage 0
    if T == "TIS" and N == "N0" and M == "M0":
        return "Stage 0"

    # Stage IA
    if T == "T1" and N == "N0" and M == "M0":
        return "Stage IA"

    # Stage IB
    if (T in ("T0", "T1")) and N == "N1MI" and M == "M0":
        return "Stage IB"

    # Stage IIA
    if (
        ((T in ("T0", "T1")) and N == "N1" and M == "M0")
        or (T == "T2" and N == "N0" and M == "M0")
    ):
        return "Stage IIA"

    # Stage IIB
    if (
        (T == "T2" and N == "N1" and M == "M0")
        or (T == "T3" and N == "N0" and M == "M0")
    ):
        return "Stage IIB"

    # Stage IIIA
    if (
        (T in ("T0", "T1", "T2") and N == "N2" and M == "M0")
        or (T == "T3" and N in ("N1", "N2") and M == "M0")
    ):
        return "Stage IIIA"

    # Stage IIIB
    if T == "T4" and M == "M0":
        return "Stage IIIB"

    # Stage IIIC
    if N == "N3" and M == "M0":
        return "Stage IIIC"

    return "Stage not mapped (breast)"


# ---------------------------------------------------------------------
# 2) Lung Cancer (NSCLC) – AJCC 8
# ---------------------------------------------------------------------

def stage_nsclc(T: str, N: str, M: str) -> str:
    # Stage IV
    if M in ("M1A", "M1B"):
        return "Stage IVA"
    if M == "M1C":
        return "Stage IVB"
    if M and M != "M0":
        # any other M1 pattern
        return "Stage IV"

    # Stage 0
    if T == "TIS" and N == "N0" and M == "M0":
        return "Stage 0"

    # Stage IA1
    if T in ("T1A", "T1A(MI)", "T1AMI") and N == "N0" and M == "M0":
        return "Stage IA1"

    # Stage IA2
    if T == "T1B" and N == "N0" and M == "M0":
        return "Stage IA2"

    # Stage IA3
    if T == "T1C" and N == "N0" and M == "M0":
        return "Stage IA3"

    # Stage IB
    if T == "T2A" and N == "N0" and M == "M0":
        return "Stage IB"

    # Stage IIA
    if T == "T2B" and N == "N0" and M == "M0":
        return "Stage IIA"

    # Stage IIB
    if (T == "T3" and N == "N0" and M == "M0") or (
        T in ("T1", "T2") and N == "N1" and M == "M0"
    ):
        return "Stage IIB"

    # Stage IIIA
    if (
        T in ("T1", "T2") and N == "N2" and M == "M0"
    ) or (T == "T3" and N in ("N1", "N2") and M == "M0"):
        return "Stage IIIA"

    # Stage IIIB
    if (
        T in ("T3", "T4") and N == "N2" and M == "M0"
    ) or (T in ("T1", "T2") and N == "N3" and M == "M0"):
        return "Stage IIIB"

    # Stage IIIC
    if T in ("T3", "T4") and N == "N3" and M == "M0":
        return "Stage IIIC"

    return "Stage not mapped (NSCLC)"


# ---------------------------------------------------------------------
# 3) Colon / Rectal Cancer – AJCC 8
# ---------------------------------------------------------------------

def stage_colorectal(T: str, N: str, M: str) -> str:
    # Stage IV
    if M == "M1A":
        return "Stage IVA"
    if M == "M1B":
        return "Stage IVB"
    if M == "M1C":
        return "Stage IVC"
    if M and M != "M0":
        return "Stage IV"

    # Stage 0
    if T == "TIS" and N == "N0" and M == "M0":
        return "Stage 0"

    # Stage I
    if T in ("T1", "T2") and N == "N0" and M == "M0":
        return "Stage I"

    # Stage II
    if T == "T3" and N == "N0" and M == "M0":
        return "Stage IIA"
    if T == "T4A" and N == "N0" and M == "M0":
        return "Stage IIB"
    if T == "T4B" and N == "N0" and M == "M0":
        return "Stage IIC"

    # Stage IIIA
    if (
        T in ("T1", "T2") and N in ("N1", "N1C") and M == "M0"
    ) or (T == "T1" and N == "N2A" and M == "M0"):
        return "Stage IIIA"

    # Stage IIIB
    if (
        (T in ("T3", "T4A") and N in ("N1", "N1C") and M == "M0")
        or (T in ("T2", "T3") and N == "N2A" and M == "M0")
        or (T in ("T1", "T2") and N == "N2B" and M == "M0")
    ):
        return "Stage IIIB"

    # Stage IIIC
    if (T in ("T4A", "T4B") and N in ("N2", "N2A", "N2B") and M == "M0"):
        return "Stage IIIC"

    return "Stage not mapped (colorectal)"


# ---------------------------------------------------------------------
# 4) Stomach (Gastric Cancer) – AJCC 8
# ---------------------------------------------------------------------

def stage_stomach(T: str, N: str, M: str) -> str:
    # Stage IV
    if M and M != "M0":
        return "Stage IV"

    # Stage I
    if (T == "T1" and N in ("N0", "N1")) or (T == "T2" and N == "N0"):
        return "Stage I"

    # Stage II
    if (T == "T1" and N == "N2") or (T == "T2" and N == "N1") or (
        T == "T3" and N == "N0"
    ):
        return "Stage II"

    # Stage III
    if (T == "T2" and N == "N2") or (T == "T3" and N in ("N1", "N2")) or (
        T == "T4A" and N in ("N0", "N1", "N2")
    ) or (T == "T4B" and N in ("N0", "N1")):
        return "Stage III"

    return "Stage not mapped (stomach)"


# ---------------------------------------------------------------------
# 5) Pancreatic Cancer – AJCC 8
# ---------------------------------------------------------------------

def stage_pancreas(T: str, N: str, M: str) -> str:
    # Stage IV
    if M and M != "M0":
        return "Stage IV"

    # Stage I
    if T in ("T1", "T2") and N == "N0" and M == "M0":
        return "Stage I"

    # Stage II
    if (T == "T3" and N == "N0" and M == "M0") or (
        T in ("T1", "T2", "T3") and N == "N1" and M == "M0"
    ):
        return "Stage II"

    # Stage III
    # mapping given: T4 (unresectable arterial invasion) OR N2 (≥4 nodes), with M0
    if (T == "T4" and M == "M0") or (N == "N2" and M == "M0"):
        return "Stage III"

    return "Stage not mapped (pancreas)"


# ---------------------------------------------------------------------
# 6) Liver Cancer (Hepatocellular Carcinoma) – AJCC 8
# ---------------------------------------------------------------------

def stage_liver(T: str, N: str, M: str) -> str:
    # Stage IV subdivided:
    if M and M != "M0":
        return "Stage IVB"
    if N == "N1" and M == "M0":
        return "Stage IVA"

    # Stage I
    if T == "T1" and N == "N0" and M == "M0":
        return "Stage I"

    # Stage II
    if T == "T2" and N == "N0" and M == "M0":
        return "Stage II"

    # Stage III
    if (T == "T3" and N == "N0" and M == "M0") or (
        T == "T4" and N == "N0" and M == "M0"
    ):
        return "Stage III"

    return "Stage not mapped (liver)"


# ---------------------------------------------------------------------
# 7) Melanoma – AJCC 8
# ---------------------------------------------------------------------

def stage_melanoma(T: str, N: str, M: str) -> str:
    # Stage 0
    if T == "TIS" and N == "N0" and M == "M0":
        return "Stage 0"

    # Stage IV
    if M in ("M1A", "M1B", "M1C", "M1D") or (M and M != "M0"):
        if M == "M1A":
            return "Stage IV (M1a)"
        if M == "M1B":
            return "Stage IV (M1b)"
        if M == "M1C":
            return "Stage IV (M1c)"
        if M == "M1D":
            return "Stage IV (M1d)"
        return "Stage IV"

    # Stage I
    if T == "T1A" and N == "N0" and M == "M0":
        return "Stage IA"
    if T in ("T1B", "T2A") and N == "N0" and M == "M0":
        return "Stage IB"

    # Stage II
    if T in ("T2B", "T3A") and N == "N0" and M == "M0":
        return "Stage IIA"
    if T in ("T3B", "T4A") and N == "N0" and M == "M0":
        return "Stage IIB"
    if T == "T4B" and N == "N0" and M == "M0":
        return "Stage IIC"

    # Stage III
    if N in ("N1", "N2", "N3") and M == "M0":
        return "Stage III"

    return "Stage not mapped (melanoma)"


# ---------------------------------------------------------------------
# 8) Prostate Cancer – Anatomic Logic
# ---------------------------------------------------------------------

def stage_prostate(T: str, N: str, M: str) -> str:
    # Stage IV
    if M and M != "M0":
        return "Stage IVB"
    if N == "N1" and M == "M0":
        return "Stage IVA"

    # Stage I
    if T in ("T1", "T2A") and N == "N0" and M == "M0":
        return "Stage I"

    # Stage II (no A/B/C separation here – depends on grade/PSA in real AJCC)
    if T in ("T2B", "T2C") and N == "N0" and M == "M0":
        return "Stage II"

    # Stage IIIA
    if T == "T3A" and N == "N0" and M == "M0":
        return "Stage IIIA"

    # Stage IIIB
    if T in ("T3B", "T4") and N == "N0" and M == "M0":
        return "Stage IIIB"

    return "Stage not mapped (prostate)"


# ---------------------------------------------------------------------
# 9) Kidney (Renal Cell Carcinoma)
# ---------------------------------------------------------------------

def stage_kidney(T: str, N: str, M: str) -> str:
    # Stage IV
    if T == "T4" and M == "M0":
        return "Stage IV"
    if M and M != "M0":
        return "Stage IV"

    # Stage I
    if T == "T1" and N == "N0" and M == "M0":
        return "Stage I"

    # Stage II
    if T == "T2" and N == "N0" and M == "M0":
        return "Stage II"

    # Stage III
    if (T == "T3" and N == "N0" and M == "M0") or (N == "N1" and M == "M0"):
        return "Stage III"

    return "Stage not mapped (kidney)"


# ---------------------------------------------------------------------
# 10) Bladder Cancer – AJCC 8
# ---------------------------------------------------------------------

def stage_bladder(T: str, N: str, M: str) -> str:
    # Stage IV
    if T == "T4B" and M == "M0":
        return "Stage IV"
    if N in ("N1", "N2", "N3") and M == "M0":
        return "Stage IV"
    if M and M != "M0":
        return "Stage IV"

    # Stage 0
    if T in ("TA", "TIS") and N == "N0" and M == "M0":
        return "Stage 0"

    # Stage I
    if T == "T1" and N == "N0" and M == "M0":
        return "Stage I"

    # Stage II
    if T == "T2" and N == "N0" and M == "M0":
        return "Stage II"

    # Stage III
    if T in ("T3", "T4A") and N == "N0" and M == "M0":
        return "Stage III"

    return "Stage not mapped (bladder)"


# ---------------------------------------------------------------------
#  FastAPI Layer
# ---------------------------------------------------------------------

app = FastAPI(
    title="AJCC TNM Staging API",
    description=(
        "Simplified AJCC anatomic staging API for multiple cancer types.\n"
        "Not a full official AJCC implementation; uses custom mapping logic."
    ),
    version="1.0.0",
)


class StagingRequest(BaseModel):
    cancer_type: str
    T: str
    N: str
    M: str


class StagingResponse(BaseModel):
    cancer_type: str
    T: str
    N: str
    M: str
    stage: str


@app.get("/health", tags=["system"])
def health_check():
    return {"status": "ok"}


@app.post("/stage", response_model=StagingResponse, tags=["staging"])
def get_stage(body: StagingRequest):
    """
    Return AJCC stage group based on cancer type and TNM.
    """
    stage = stage_cancer(
        cancer_type=body.cancer_type,
        T=body.T,
        N=body.N,
        M=body.M,
    )

    return StagingResponse(
        cancer_type=body.cancer_type,
        T=body.T,
        N=body.N,
        M=body.M,
        stage=stage,
    )


if __name__ == "__main__":
    import uvicorn
    uvicorn.run("ajcc_api:app", host="0.0.0.0", port=8000, reload=True)