from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from huggingface_hub import hf_hub_download
from llama_cpp import Llama
import json
import json_repair

# --- 1. CONFIGURATION ---
print(">>> INITIALIZING SomAI TEXT NODE...")
# Text Model Only
PHI3_REPO = "microsoft/Phi-3-mini-4k-instruct-gguf"
PHI3_FILE = "Phi-3-mini-4k-instruct-q4.gguf"

app = FastAPI()
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_headers=["*"],
)

# --- 2. ICD-10 VALIDATION DATABASE (MOCK) ---
VALID_ICD10 = {
    "I10": "Essential (primary) hypertension",
    "E11.9": "Type 2 diabetes mellitus without complications",
    "E78.5": "Hyperlipidemia, unspecified",
    "Z86.73": "Personal history of transient ischemic attack (TIA), and cerebral infarction without residual deficits",
    "R07.9": "Chest pain, unspecified",
    "J45.909": "Unspecified asthma, uncomplicated",
    "F41.1": "Generalized anxiety disorder",
    "G44.209": "Tension-type headache, unspecified, not intractable",
    "R69": "Illness, unspecified"
}

def validate_codes(pipeline):
    """Checks generated codes against allowlist and flags suspicious ones."""
    validated = []
    for item in pipeline:
        code = item.get('code', '').upper().replace('.', '')
        match = next((k for k in VALID_ICD10 if k.replace('.', '') == code), None)

        if match:
            validated.append({
                "type": item.get("type", "Unknown"),
                "code": match,
                "description": VALID_ICD10[match] + " [VERIFIED]"
            })
        else:
            validated.append({
                "type": item.get("type", "Unknown"),
                "code": "R69",
                "description": f"Unverified Code ({item.get('code')}) - Clinical Review Required"
            })
    return validated

# --- 3. MODEL LOADERS ---
llm = None
try:
    print(">>> LOADING TEXT ENGINE...")
    model_path = hf_hub_download(repo_id=PHI3_REPO, filename=PHI3_FILE)
    llm = Llama(model_path=model_path, n_ctx=2048, n_threads=2, verbose=False)
except Exception as e:
    print(f"!!! TEXT ENGINE FAILED: {e}")

# --- 4. REQUEST MODELS ---
class AnalysisRequest(BaseModel):
    age: int
    condition: str
    history: str | None = None
    allergies: str | None = None
    systolicBp: int | None = None
    systolic_bp: int | None = None
    glucose: int
    heartRate: int | None = None
    spo2: int | None = None
    weight: float | None = None
    temperature: float | None = None
    sleepQuality: int | None = None
    missedDoses: int | None = None
    clinicalNote: str | None = None
    riskScore: int | None = None
    prompt: str | None = None

class ChatRequest(BaseModel):
    prompt: str

# --- 5. ROUTES ---
@app.get("/")
def health():
    return {"status": "Active", "system": "SomAI Text Node"}

@app.post("/analyze")
def analyze(req: AnalysisRequest):
    if not llm: raise HTTPException(503, "Text Model Unavailable")

    formatted_prompt = f"<|user|>{req.prompt}<|end|><|assistant|>"
    output = llm(formatted_prompt, max_tokens=400, temperature=0.1, stop=["<|end|>"])
    raw_text = output['choices'][0]['text']

    try:
        data = json_repair.repair_json(raw_text, return_objects=True)
    except:
        data = {
            "summary": "Manual review recommended.", 
            "actionItems": ["Consult Physician"], 
            "insuranceNote": "Review required.",
            "primaryConditionCode": {"code": "R69", "description": "Unspecified"},
            "historyCodes": []
        }

    p_code = data.get('primaryConditionCode', {})
    validated_p = validate_codes([{"code": p_code.get('code'), "type": "Primary", "description": p_code.get('description')}])

    h_codes = data.get('historyCodes', [])
    validated_h = validate_codes([{"code": h.get('code'), "type": "History", "description": h.get('description')} for h in h_codes])

    final_response = {
        "summary": data.get("summary"),
        "actionItems": data.get("actionItems"),
        "primaryConditionCode": validated_p[0],
        "historyCodes": validated_h,
        "insuranceNote": data.get("insuranceNote") + " [Validated by SomAI Engine]"
    }
    return final_response

@app.post("/generate")
def generate(req: ChatRequest):
    if not llm: raise HTTPException(503, "Text Model Unavailable")

    formatted_prompt = f"<|user|>{req.prompt}<|end|><|assistant|>"
    output = llm(formatted_prompt, max_tokens=250, temperature=0.6, stop=["<|end|>"])
    return {"text": output['choices'][0]['text'].strip()}