import json
import zipfile
import io
import requests
import os
import boto3
import csv

def lambda_handler(event, context):
    # Gets the Well Data from the Open Source Repo
    s3 = boto3.resource('s3')
    with requests.get('https://www.data.bsee.gov/Well/Files/eWellWARRawData.zip', stream=True) as r:
        z = zipfile.ZipFile(io.BytesIO(r.content))
        os.chdir('/tmp/')
        
        z.extractall()
        print(os.listdir())
        os.chdir('./eWellWARRawData/')

    count = 0
    lines2save=[]
    file_number = 1
    first = True
    headers = ''
    with open('mv_war_main.txt') as infile:
        for line in infile:
            # Some special logic to handle the first line, saves the headers for each one
            if first:
                headers = line
                lines2save.append(headers)
                count +=1
                first = False
                continue
            
            # Otherwise continue on to the rest of the loop
            lines2save.append(line)
            if count%5000==0:
                file_name = 'mv_war_main_'+str(file_number)+'.txt'
                outfiles = open(file_name,"w")
                outfiles.writelines(lines2save)
                outfiles.close()
                lines2save = []
                lines2save.append(headers)
                s3.meta.client.upload_file('/tmp/eWellWARRawData/'+file_name, os.environ['S3Bucket'], 'mv_war_main/'+file_name)
                file_number+=1
                os.remove(file_name)
            count +=1
        file_name = 'mv_war_main_'+str(file_number)+'.txt'
        outfiles = open(file_name,"w")
        outfiles.writelines(lines2save)
        outfiles.close()
        s3.meta.client.upload_file('/tmp/eWellWARRawData/'+file_name, os.environ['S3Bucket'], 'mv_war_main/'+file_name)

    # Do it again for mv_war_main_prop
    count = 0
    lines2save=[]
    file_number = 1
    first = True
    headers = ''
    with open('mv_war_main_prop.txt') as infile:
        for line in infile:
            # Some special logic to handle the first line, saves the headers for each one
            if first:
                headers = line
                lines2save.append(headers)
                count +=1
                first = False
                continue
            
            # Otherwise continue on to the rest of the loop
            lines2save.append(line)
            if count%5000==0:
                file_name = 'mv_war_main_prop_'+str(file_number)+'.txt'
                outfiles = open(file_name,"w")
                outfiles.writelines(lines2save)
                outfiles.close()
                lines2save = []
                lines2save.append(headers)
                s3.meta.client.upload_file('/tmp/eWellWARRawData/'+file_name, os.environ['S3Bucket'], 'mv_war_main_prop/'+file_name)
                file_number+=1
                os.remove(file_name)
            count +=1
        file_name = 'mv_war_main_prop_'+str(file_number)+'.txt'
        outfiles = open(file_name,"w")
        outfiles.writelines(lines2save)
        outfiles.close()
        s3.meta.client.upload_file('/tmp/eWellWARRawData/'+file_name, os.environ['S3Bucket'], 'mv_war_main_prop/'+file_name)

    # All that's remaining is to upload the remarks data to S3
    s3.meta.client.upload_file('/tmp/eWellWARRawData/mv_war_main_prop_remark.txt', os.environ['S3Bucket'], 'inputs/mv_war_main_prop_remark.txt')
    return ({'UploadBSEE':'Success'})