import os
import pandas as pd
from sodapy import Socrata
from schemas import MetaData, FieldStats
from google.cloud import bigquery
from google.cloud.exceptions import NotFound

class MyData:
    #bigquery_table: socrata_dataset_id
    id_dictionary = {
        "parking_lax": "dik5-hwp6",
        "parking_dot": "e7h6-4a3e",
        # "myla311_2018": "h65r-yf5i",
        # "myla311_2019": "pvft-t768",
        # "myla311_2020": "rq3b-xjk8",
        # "myla311_2021": "97z7-y5bt",
        # "myla311_2022": "i5ke-k6by"
    }

    on_dictionary = {
        "parking_lax": ["key_value", "dataexportdatetime"],
        "parking_dot": ["spaceid", "eventtime"],
        "myla311_2018": ["srnumber"],
        "myla311_2019": ["srnumber"],
        "myla311_2020": ["srnumber"],
        "myla311_2021": ["srnumber"],
        "myla311_2022": ["srnumber"]
    }

    max_dictionary ={
        "parking_lax": "dataexportdatetime",
        "parking_dot": "eventtime",
        "myla311_2018": "updateddate",
        "myla311_2019": "updateddate",
        "myla311_2020": "updateddate",
        "myla311_2021": "updateddate",
        "myla311_2022": "updateddate"
    }





class MyAppConfig:
    SOCRATA_APP_TOKEN=os.environ['SOCRATA_APP_TOKEN']
    SOCRATA_APP_KEY=os.environ['SOCRATA_APP_KEY']
    SOCRATA_APP_SECRET=os.environ['SOCRATA_APP_SECRET']
    SOCRATA_APP_DOMAIN=os.environ['SOCRATA_APP_DOMAIN']
    GOOGLE_PROJECT=os.environ['GOOGLE_PROJECT']
    GOOGLE_DATASET=os.environ['GOOGLE_DATASET']


class MySocrata(MyAppConfig):
    def __init__(self) -> None:
        super().__init__()
        self.set_client()
    
    def set_client(self) -> None:
        self.client = Socrata(
            self.SOCRATA_APP_DOMAIN,
            self.SOCRATA_APP_TOKEN,
            username=self.SOCRATA_APP_KEY,
            password=self.SOCRATA_APP_SECRET
        )

    def set_dataset_id(self, dataset_id) -> None:
        self.dataset_id = dataset_id

    def set_metadata(self) -> None:
        response = self.client.get_metadata(self.dataset_id)
        self.metadata = response


    def set_schema(self) -> None:
        self.schema = MetaData(**self.metadata)

    
    def set_dataframe(self, limit: int, where: str):
        response = self.client.get(
            self.dataset_id, 
            limit=limit, 
            where=where
        )

        df = pd.DataFrame.from_records(response)
        for column in self.schema.columns:
            if column not in df.columns:
                df[column] = None

        df = df[self.schema.columns]

        for column in self.schema.columndates:
            df[column] = pd.to_datetime(df[column], format='%Y-%m-%dT%H:%M:%S', errors='coerce', utc=False )

        vectors = [self.schema.columnnumbers, self.schema.columntexts, self.schema.columnothers]
        astypes = ['float64', 'string', 'string']

        for i, vector in enumerate(vectors):
            df[vector] = df[vector].astype(astypes[i])

        self.dataframe = df

    def try_astype_Int64(self):
        df = self.dataframe      
        for column in self.schema.columnnumbers:
            try:
                df[column] = df[column].astype('float').astype('Int64')
            except TypeError:
                pass
        self.dataframe = df


class MyBigQuery(MyAppConfig):
    def __init__(self) -> None:
        super().__init__()
        self.dataset_ref = f"{self.GOOGLE_PROJECT}.{self.GOOGLE_DATASET}"
        self.set_client()
    

    def set_client(self) -> None:
        self.client = bigquery.Client()
    

    def upload_dataframe(self, dataframe: pd.DataFrame, table: str, schema: list, truncate: bool):
        if truncate:
            write_disposition="WRITE_TRUNCATE"
        else:
            write_disposition="WRITE_APPEND"

        job_config = bigquery.LoadJobConfig(
            schema=schema,
            write_disposition=write_disposition
        )

        job = self.client.load_table_from_dataframe(
            dataframe, 
            f"{self.dataset_ref}.{table}",
            job_config=job_config
        )

        result = job.result()
        return result
    

    def get_sql_job(self, sql: str):
        job = self.client.query(sql)
        result = job.result() # result is an iterator
        return result


    def run_sql_job(self, sql: str) -> None:
        # does not wait for job to finish
        job = self.client.query(sql)
    
    
    def set_fieldstats(self, table:str, field:str) -> None:
        sql = (
            f"SELECT max({field}) as fmax, min({field}) as fmin, count({field}) as fcnt, count(*) as frows"
            f" FROM {self.dataset_ref}.{table}"
            f" LIMIT 1"
        )
        result = self.get_sql_job(sql)
        for row in result:
            self.fieldstats = FieldStats(fmin=row['fmin'], fmax=row['fmax'], fcnt=row['fcnt'], frows=row['frows'])
            return None


    def set_tablefields(self, table: str) -> None:
        info = self.client.get_table(f"{self.dataset_ref}.{table}")
        tablefields = {field.name:field.field_type for field in info.schema}
        self.tablefields = tablefields

    
    def set_jobconfig_schema(self, tablefields: dict) -> None:
        self.jobconfig_schema = [bigquery.SchemaField(key, tablefields[key]) for key in tablefields]


    def merge_target_source_on(self, target:str, source: str, on: list):
        fields = self.tablefields.keys()
        onsql = 'ON ' +' AND '.join([f"T.{item} = S.{item}" for item in on])
        insert = "INSERT (" + ', '.join(fields) +") VALUES (" + ', '.join(fields) + ")"
        update = "UPDATE SET " + ', '.join([f"{key} = S.{key}" for key in fields])

        sql = (
            f" MERGE `{self.dataset_ref}.{target}` T"
            f" USING `{self.dataset_ref}.{source}` S"
            f" {onsql}"
            f" WHEN MATCHED THEN"
            f" {update}"
            f" WHEN NOT MATCHED THEN"
            f" {insert}"
        )

        self.run_sql_job(sql)



    def select_on_fieldmax(self, table, on, fieldmax):
        sql = (
            f" SELECT s1.*"
            f" FROM `{self.dataset_ref}.{table} s1"
            f" INNER JOIN" 
            f" ("
            f"  SELECT {on}, max({fieldmax}) {fieldmax}"
            f"  FROM `{self.dataset_ref}.{table}`" 
            f"  GROUP BY {on}"
            f" ) s2 on s2.{on} = s1.{on} and s2.{fieldmax} = s1. {fieldmax}"
        )
        return sql
        

# MERGE dataset.Inventory T
# USING dataset.NewArrivals S
# ON T.product = S.product
# WHEN MATCHED THEN
#   UPDATE SET quantity = T.quantity + S.quantity
# WHEN NOT MATCHED THEN
#   INSERT (product, quantity) VALUES(product, quantity)



    def check_exists_table(self, table) -> bool:
        try:
            self.client.get_table(f"{self.dataset_ref}.{table}")
        except NotFound:
            return False

        return True

    def check_exists_dataset(self) -> bool:
        try:
            self.client.get_dataset(f"{self.dataset_ref}")
        except NotFound:
            return False

        return True