from __future__ import annotations
import json, urllib.parse, xml.etree.ElementTree as ET
from .common import fetch_json, fetch_bytes, cache_payload, normalized_source, clean_text

BASE="https://dailymed.nlm.nih.gov/dailymed/services/v2"
NS={'h':'urn:hl7-org:v3'}

def _section_text(section):
    chunks=[]
    for el in section.iter():
        if el.text and el.tag.split('}')[-1] not in ('title',): chunks.append(el.text)
    return clean_text(' '.join(chunks))

def research(name: str, cache_dir):
    url=f"{BASE}/spls.json?drug_name={urllib.parse.quote(name)}&pagesize=5&page=1"
    data, meta=fetch_json(url)
    cm=cache_payload(cache_dir,"dailymed_index",name,url,json.dumps(data).encode(),meta)
    rows=data.get("data") or data.get("DATA") or []
    # v2 may return list of dicts; support v1-style arrays too
    spls=[]
    if isinstance(rows,list):
        for row in rows:
            if isinstance(row,dict): spls.append(row)
            elif isinstance(row,list) and row: spls.append({"setid":row[0],"title":row[1] if len(row)>1 else name,"spl_version":row[2] if len(row)>2 else "","published_date":row[3] if len(row)>3 else ""})
    sources=[]
    for row in spls[:3]:
        setid=row.get("setid") or row.get("SETID") or row.get("set_id") or ""
        if not setid: continue
        xml_url=f"{BASE}/spls/{setid}.xml"
        body,xmeta=fetch_bytes(xml_url,headers={"Accept":"application/xml,text/xml"})
        xcm=cache_payload(cache_dir,"dailymed_spl",setid,xml_url,body,xmeta)
        sections=[]
        try:
            root=ET.fromstring(body)
            for sec in root.findall('.//h:section',NS):
                title_el=sec.find('./h:title',NS) or sec.find('.//h:title',NS)
                title=clean_text(title_el.text if title_el is not None else "Untitled section")
                code=sec.find('./h:code',NS)
                sections.append({"section":title,"subsection":"","locator":f"SPL section code {code.get('code','') if code is not None else ''}","text":_section_text(sec),"sectionCode":code.get('code','') if code is not None else ''})
        except ET.ParseError as exc:
            sections=[{"section":"Parse error","subsection":"","locator":"XML","text":str(exc)}]
        title=row.get("title") or row.get("TITLE") or f"DailyMed SPL {setid}"
        sources.append(normalized_source(connector="dailymed",title=title,url=xml_url,source_type="structured_product_label",
            version=str(row.get("spl_version") or row.get("SPL_VERSION") or ""), publication_date=str(row.get("published_date") or row.get("PUBLISHED_DATE") or ""),
            sections=sections,raw_meta=xcm,identifiers={"setId":setid},warnings=[]))
    return sources
