from typing import Annotated
from fastapi import FastAPI, Query
from fastapi.responses import JSONResponse
import xml.etree.ElementTree as xml
import requests
import json

app = FastAPI()

@app.get("/api/v1/health")
def health():
    return { "status": "ok" }

@app.get("/api/v1/lookup/{domain}")
def read_root(
    domain: str,
    platform: Annotated[
        str | None,
        Query(pattern="^default|gcc|gcc_high$")
    ] = "default",
):

    match platform:
        case "gcc_high": base_uri = "autodiscover-s.office365.us"
        case "gcc": base_uri = "autodiscover-s-dod.office365.us"
        case "default": base_uri = "autodiscover-s.outlook.com"

    headers = { "Content-Type": "text/xml; charset=utf-8", "SOAPAction": "http://schemas.microsoft.com/exchange/2010/Autodiscover/Autodiscover/GetFederationInformation" }
    uri = f"https://{base_uri}/autodiscover/autodiscover.svc"
    data = f"""<?xml version="1.0" encoding="utf-8"?>
    <soap:Envelope xmlns:exm="http://schemas.microsoft.com/exchange/services/2006/messages" xmlns:ext="http://schemas.microsoft.com/exchange/services/2006/types" xmlns:a="http://www.w3.org/2005/08/addressing" xmlns:soap="http://schemas.xmlsoap.org/soap/envelope/" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xmlns:xsd="http://www.w3.org/2001/XMLSchema">
        <soap:Header>
            <a:Action soap:mustUnderstand="1">http://schemas.microsoft.com/exchange/2010/Autodiscover/Autodiscover/GetFederationInformation</a:Action>
            <a:To soap:mustUnderstand="1">{uri}</a:To>
            <a:ReplyTo><a:Address>http://www.w3.org/2005/08/addressing/anonymous</a:Address></a:ReplyTo>
        </soap:Header>
        <soap:Body>
            <GetFederationInformationRequestMessage xmlns="http://schemas.microsoft.com/exchange/2010/Autodiscover">
                <Request>
                    <Domain>{domain}</Domain>
                </Request>
            </GetFederationInformationRequestMessage>
        </soap:Body>
    </soap:Envelope>
    """

    response = requests.post(uri, data=data, headers=headers)
    if (not response.ok): return { "error": f"Error contacting autodiscover service for {domain}" }

    root = xml.fromstring(response.text)
    content = [i.text for i in root[1][0][0][3]] # root/body/GetFederationInformationResponseMessage/Domains

    return JSONResponse(content=content, headers={"Cache-Control": "max-age=604800, public"})