import re
from fastapi import FastAPI
from pydantic import BaseModel

app = FastAPI()

class StudyText(BaseModel):
    text: str

def extract_study_timeline(text: str):
    screening = re.search(r'(?:Screening|Pre-study observation|Initial Check)(?:\s*(?:phase|period))?\s*(?:is|:|-|of|lasts)?\s*(\d+)\s*weeks?', text, re.IGNORECASE)
    treatment = re.search(r'(?:Treatment|Intervention|Therapy|Dosing phase|Main study(?:\s*period)?)(?:\s*(?:is|:|-|of|lasts))?\s*(\d+)\s*weeks?', text, re.IGNORECASE)
    follow_up = re.search(r'(?:Follow[-\s]*up|Recovery|Post-study monitoring|Observation phase|After-treatment)(?:\s*(?:is|:|-|of|lasts))?\s*(\d+)\s*weeks?', text, re.IGNORECASE)

    return {
        "Screening": int(screening.group(1)) if screening else None,
        "Treatment": int(treatment.group(1)) if treatment else None,
        "Follow-Up": int(follow_up.group(1)) if follow_up else None
    }

@app.post("/extract-timeline/")
async def extract_timeline(request: StudyText):
    return extract_study_timeline(request.text)

@app.get("/")
async def root():
    return {"message": "Study Timeline Extractor API is running. Use /extract-timeline/ to extract data."}