import numpy as np
import json
import pandas as pd

import requests
import requests_cache

requests_cache.install_cache('cache', allowable_methods=("POST","GET"))

def get_json(restarttoken=None, count=3000, active=True):
    if not restarttoken:
        restarttoken = ""

    else:
        restarttoken = json.dumps(restarttoken)
        print(restarttoken)

    modelid = 404287
    if not active:
        modelid = 404284

    headers = {
        'Connection': 'keep-alive',
        'sec-ch-ua': '"Chromium";v="88", "Google Chrome";v="88", ";Not A Brand";v="99"',
        'DNT': '1',
        'sec-ch-ua-mobile': '?0',
        'User-Agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_14_0) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/88.0.4324.192 Safari/537.36',

        # active
        'ActivityId': '66b5ae32-554c-470a-9c78-b54a8718b77c',
        'Accept': 'application/json, text/plain, */*',
        'RequestId': '2c7ab566-5a12-fe1a-0ac5-c54dd3d9c532',
        'X-PowerBI-ResourceKey': 'b2c8d2f2-3ad1-48dc-883c-d4163a6e2d8f',


        'Content-Type': 'application/json;charset=UTF-8',
        'Origin': 'https://app.powerbigov.us',
        'Sec-Fetch-Site': 'cross-site',
        'Sec-Fetch-Mode': 'cors',
        'Sec-Fetch-Dest': 'empty',
        'Referer': 'https://app.powerbigov.us/',
        'Accept-Language': 'en-US,en;q=0.9',
    }

    if not active:
        headers.update({
            # inactive
            'ActivityId': '4e140a40-456c-4847-bd80-f17e1ddaf836',
            'Accept': 'application/json, text/plain, */*',
            'RequestId': '2814cc0b-5dce-7363-7dba-b948e9387405',
            'X-PowerBI-ResourceKey': '87914378-578f-4f43-b75e-8ddaeafbdda2',
        })


    params = (
        ('synchronous', 'true'),
    )


    data = """
    {
      "version": "1.0.0",
      "queries": [
        {
          "Query": {
            "Commands": [
              {
                "SemanticQueryDataShapeCommand": {
                  "Query": {
                    "Version": 2,
                    "From": [
                      {
                        "Name": "q1",
                        "Entity": "CCRB Active - Oracle",
                        "Type": 0
                      }
                    ],
                    "Select": [


    { "Name": "Query1.Unique Id", "Column": { "Expression": { "SourceRef": { "Source": "q1" } }, "Property": "Unique Id" }, },
    { "Name": "Query1.Command1", "Column": { "Expression": { "SourceRef": { "Source": "q1" } }, "Property": "Command" }, },
    { "Name": "Query1.Last Name1", "Column": { "Expression": { "SourceRef": { "Source": "q1" } }, "Property": "Last Name" }, },
    { "Name": "Query1.First Name1", "Column": { "Expression": { "SourceRef": { "Source": "q1" } }, "Property": "First Name" }, },
    { "Name": "Query1.Rank1", "Column": { "Expression": { "SourceRef": { "Source": "q1" } }, "Property": "Rank" }, },
    { "Name": "Query1.ShieldNo", "Column": { "Expression": { "SourceRef": { "Source": "q1" } }, "Property": "Shield No" }, },
    { "Name": "Query1.Complaint Id", "Column": { "Expression": { "SourceRef": { "Source": "q1" } }, "Property": "Complaint ID" }, },
    { "Name": "Query1.Sub Complaints", "Column": { "Expression": { "SourceRef": { "Source": "q1" } }, "Property": "Sub Complaints" }, },
    { "Name": "Query1.Total Complaints", "Column": { "Expression": { "SourceRef": { "Source": "q1" } }, "Property": "Total Complaints" }, },



                    ],
                    "Where": [


                    ],
                    "OrderBy": [
                      {
                        "Direction": 1,
                        "Expression": {
                          "Column": {
                            "Expression": {
                              "SourceRef": {
                                "Source": "q1"
                              }
                            },
                            "Property": "Command"
                          }
                        }
                      }
                    ]
                  },
                  "Binding": {
                    "Primary": {
                      "Groupings": [
                        {
                          "Projections": [
                            0,
                            1,
                            2,
                            3,
                            4,
                            5,
                            6,
                            7,
                            8,
                          ]
                        }
                      ]
                    },
                    "DataReduction": {
                      "Primary": {
                        "Window": {
                          "Count": %i,
                          "RestartTokens": [%s],
                        }
                      }
                    },
                    "Version": 1
                  }
                }
              }
            ]
          },
          "QueryId": "",
          "ApplicationContext": {
            "DatasetId": "523ab509-8e2d-43ed-bfad-11fcd05180d7",
            "Sources": [
              {
                "ReportId": "f508555a-b39d-4c10-8d46-a14bc282e079"
              }
            ]
          }
        }
      ],
      "cancelQueries": [],
      "modelId": %i
    }
    """ % (count, restarttoken, modelid)

    response = requests.post('https://wabi-us-gov-virginia-api.analysis.usgovcloudapi.net/public/reports/querydata', headers=headers, params=params, data=data)

    print(response.from_cache)

    return response.json()


def parse_json(data):
    if "results" not in data:
        print(data)
    try:
        restarttoken = data["results"][0]["result"]["data"]["dsr"]["DS"][0]["RT"][0]
    except:
        restarttoken = None
    # print(restarttoken)
    rawrows = data["results"][0]["result"]["data"]["dsr"]["DS"][0]["PH"][0]["DM0"]

    rawlookup = data["results"][0]["result"]["data"]["dsr"]["DS"][0]["ValueDicts"]
    lookup = dict()
    for k, v in rawlookup.items():
        lookup[int(k[1:])] = v
        # print(k, len(v), v)
        # print()

    def get_indices(mask):
        indices = []
        for i in range(16):
            if mask & (2 << i):
                indices.append(i+1)
        return indices

    rows = []
    # allowed_lookups = set()
    for ix,x in enumerate(rawrows):
        # if ix == 0:
        #     print(x)

        rawrow = x["C"]
        mask = x.get("R", 0)
        indices = []
        if mask > 0:
            indices = get_indices(mask)
            for idx in indices:
                rawrow.insert(idx, None)
        row = []
        for i,c in enumerate(rawrow):
            if isinstance(c, int) and (i in lookup) and (c < 1e6):
                row.append(lookup[i][c])
            else:
                row.append(c)

        rows.append(row)

    df = pd.DataFrame(rows)
    df = df.ffill()
    return restarttoken, df

if __name__ == "__main__":

    count = 10000

    dfs = []

    for active in [True, False]:
        restarttoken = None
        for i in range(50):
            if i > 0 and not restarttoken:
                break
            data = get_json(restarttoken, count, active=active)
            restarttoken, df = parse_json(data)
            print(active, i, restarttoken, df.shape)
            df["active"] = active
            dfs.append(df)


    dfc = pd.concat(dfs)
    dfc.columns = [
            "id",
            "command",
            "lastname",
            "firstname",
            "rank",
            "shieldno",
            "complaintid",
            "subcomplaints",
            "totcomplaints",
            "active",
            ]
    print(dfc)

    dfc = dfc.dropna()
    dfc["totcomplaints"] = dfc["totcomplaints"].astype(int)
    dfc = dfc.drop_duplicates(["command","lastname","firstname","shieldno"]).drop(columns=["complaintid"])
    dfc = dfc.reset_index(drop=True)

    dfc.to_json("overview.jsonl.gz", lines=True, orient="records")