WITH all_file_revisions AS (
        SELECT wiki_id, revision_id, page_id, page_title, page_redirect_target, 
        revision_content_slots['main'].content_body AS revision_text,
        ROW_NUMBER() OVER (PARTITION BY wiki_id,page_id ORDER BY revision_dt DESC) AS revision_number_desc
        FROM wmf_content.mediawiki_content_history_v1
        WHERE page_namespace_id=0
    )
    SELECT wiki_id, revision_id, page_id, page_title, page_redirect_target, revision_text
    FROM all_file_revisions WHERE revision_number_desc=1

SCHEMA = StructType(
    [
        StructField('depicts', ArrayType(StringType()), False),
    ]
)

@udf(returnType=SCHEMA)
def parse_mediainfo(mediainfo: str) -> list:

    depicted = []

    try:
        mediainfo_dict = json.loads(mediainfo)
        if mediainfo_dict["statements"].get('P180'):
            snaks = mediainfo_dict["statements"].get('P180')
            for snak in snaks:
                depicted.append(snak['mainsnak']['datavalue']['value']['id'])

        if mediainfo_dict["statements"].get('P6243'):
            snaks = mediainfo_dict["statements"].get('P6243')
            for snak in snaks:
                depicted.append(snak['mainsnak']['datavalue']['value']['id'])

    except Exception as e:
        # do nothing, just allow the empty array to be returned
        pass
    
    return { "depicts": sorted(set(depicted)) }

def get_depicts(spark: SparkSession) -> DataFrame:
    return spark.sql("""
        WITH 
        commons_mediainfo AS (
            SELECT
                page_id,
                page_title,
                revision_content_slots['mediainfo'].content_body AS mediainfo,
                revision_dt,
                ROW_NUMBER() OVER (PARTITION BY wiki_id,page_id ORDER BY revision_dt DESC) AS revision_number_desc
            FROM wmf_content.mediawiki_content_history_v1
            WHERE wiki_id="commonswiki"
            AND page_namespace_id=6
        ),
        latest_rev_with_depicts_array AS (
            select page_id,parse_mediainfo(mediainfo).depicts as depicts from commons_mediainfo where revision_number_desc=1
        )
        select page_id, item_id
        from latest_rev_with_depicts_array
        lateral view explode(depicts) as item_id
    """)..withColumn('depicts', F.lit(1))

import urllib.parse
import json
from pyspark.sql.functions import udf


SCHEMA = StructType(
    [
        StructField('links', ArrayType(MapType(StringType(), StringType())), True),
        StructField('statements', MapType(StringType(), ArrayType(StringType())), True),
    ]
)

@udf(returnType=SCHEMA)
def parse_wikidata(wikidata_json: str) -> dict:

    links = []
    statements = {
        "P18": [],
        "P31": [],
        "P373": [],
    }

    try:
        wikidata_dict = json.loads(wikidata_json)
        sitelinks = wikidata_dict["sitelinks"]
        for wiki_db, sitelink in sitelinks.items():    
            links.append({"wiki_db":sitelink["site"], "page_title":urllib.parse.unquote(sitelink["title"])})
        claims = wikidata_dict["claims"]
        for property_id, property_claims in claims.items():
            if property_id == "P31":
                for claim in property_claims:
                    statements[property_id].append(claim["mainsnak"]["datavalue"]["value"]["id"])
            elif property_id == "P18" or property_id == "P373":
                for claim in property_claims:
                    statements[property_id].append(claim["mainsnak"]["datavalue"]["value"])
                
    except json.JSONDecodeError as e:
        # wikidata json can't be decoded - do nothing, just allow empty data to be returned
        pass
    
    return {"links":links, "statements": statements}
    
spark.udf.register('parse_wikidata', parse_wikidata)

def load_parsed_wikidata(spark)
    return spark.sql("""
        WITH 
        wikidata AS (
            SELECT
                page_id,
                page_title,
                revision_content_slots['main'].content_body AS item_json,
                ROW_NUMBER() OVER (PARTITION BY wiki_id,page_id ORDER BY revision_dt DESC, revision_id DESC) AS revision_number_desc
            FROM wmf_content.mediawiki_content_history_v1
            WHERE wiki_id="wikidatawiki"
            AND page_namespace_id=0
        ),
        wikidata_latest AS (
            select page_id,page_title,item_json from wikidata where revision_number_desc=1
        ),
        wikidata_parsed AS (
            select page_title as item_id,parse_wikidata(item_json) as parsed from wikidata_latest
        )
        select item_id, parsed.links as links, parsed.statements.P18 as p18, parsed.statements.P31 as p31, parsed.statements.P373 as p373 from wikidata_parsed
    """).cache()