import requests
import pandas as pd
import os
import zipfile
from io import BytesIO
import re
import shutil
from utils.utils import Paths

def download_scot_census_data():
    # Data directories
    DATA_DIR = Paths.SCOT
    CSV_DIR = DATA_DIR / "topic_tables" / "csv"
    PARQUET_DIR = DATA_DIR / "topic_tables" / "parquet"
    # Create directories if they don't exist
    CSV_DIR.mkdir(parents=True, exist_ok=True)
    PARQUET_DIR.mkdir(parents=True, exist_ok=True)


    url = "https://www.scotlandscensus.gov.uk/media/zz85kfinmf97whklasd98gfkadft5hj4f_Topic2H_20241120_1747/Census-2022-Output-Area-v1.zip"
    # Download the zip file
    r = requests.get(url)
    if r.status_code != 200:
        raise Exception(f"Failed to download Scottish data file: {r.status_code}")
    z = zipfile.ZipFile(BytesIO(r.content))
    z.extractall("./data/tmp")

    # Create metadata table
    meta_data_table = pd.DataFrame(
        columns=[
            "Table_Name",
            "Table_ID",
            "Variable_Name",
            "Type",
        ]
    )

    oa_files = os.listdir("./data/tmp")
    for file in oa_files:
        t_tab_loc = file
        t_id = t_tab_loc.split(" - ")[0]
        if t_id == "Uv211b":
            t_id = "UV211b"
        t_name = t_tab_loc.split(" - ", 1)[1].split(".")[0]  # Remove the code but keep full name (including household when applicable)

        # Load the csv file raw (not df) and print the first 5 lines
        with open(f"./data/tmp/{t_tab_loc}", "r") as f:
            # Remove any complete empty lines
            lines = [line for line in f if line.strip() != ""]
            # Remove any lines that are all commas
            lines = [line for line in lines if not re.match(r"^,+$", line)]
            # Third is who the table includes
            table_includes = lines[2].strip()

            # Find the unit of measure
            unit = "-"
            if "Household Reference Persons " in table_includes:
                unit = "Household Reference Person"
            elif "people" in table_includes or "Persons" in table_includes:
                unit = "Person"
            elif "household" in table_includes and "people" not in table_includes:
                unit = "Household"
            else:
                print(
                    f"Unit of measure not found for {t_id}, table includes: {table_includes}"
                )

            # Remove the first 3 lines
            lines = lines[3:]

            # Find the header indices, they are any remaining lines that have a blank space before the first comma
            header_indices = [
                i for i, line in enumerate(lines) if line.strip().startswith(",")
            ]

            # Load to df
            df = pd.read_csv(
                BytesIO("".join(lines).encode()),
                index_col=0,
                header=header_indices,
                skipfooter=5,
                engine="python",
                na_values=["-"],
            )

            # Remove any rows that don't have an OA of the form S12345678
            df = df[df.index.str.match(r"S\d{8}")]

            #set the counts to int
            df = df.astype("Int64")
            # Append multi-index columns to a single column if necessary
            if df.columns.nlevels > 1:
                df.columns = [": ".join(col).strip() for col in df.columns.values]

            # If any column contains Mar-15 replace it with 3 - 15
            df.columns = df.columns.str.replace("Mar-15", "3 - 15")
            df.index.name = "OA"

            # Fill NA with 0
            df.fillna(0, inplace=True)

            if len(df) != 46363:
                # Skip the table if it doesn't have 46363 rows
                print(f"Table {t_id} has {len(df)} rows, not 46363")
                continue

            # Validate data and types
            non_numeric = df.map(lambda x: not isinstance(x, (int, float)))
            if non_numeric.any().any():
                print("Non-numeric data found in the following cells:")
                print(df[non_numeric])

            # Create new column names with zero padding
            variable_names = df.columns
            var_ids = [
                f"{t_id}{str(i).zfill(4)}" for i in range(1, len(variable_names) + 1)
            ]
            df.columns = var_ids

            # Save to new CSV and Parquet
            df.to_csv(CSV_DIR / f"{t_id}.csv")
            df.to_parquet(PARQUET_DIR / f"{t_id}.parquet")

            # Add to metadata table
            meta_data_table = pd.concat(
                [
                    meta_data_table,
                    pd.DataFrame(
                        {
                            "Variable_Name": variable_names,
                            "Variable_ID": var_ids,
                            "Table_ID": [t_id] * len(variable_names),
                            "Table_Name": [t_name] * len(variable_names),
                            "Unit": [unit] * len(variable_names),
                        }
                    ),
                ]
            )

    # Remove the temporary files and directory
    shutil.rmtree("./data/tmp")

    # Create full name column in metadata table
    meta_data_table["Full_Name"] = (
        meta_data_table["Table_Name"] + " - " + meta_data_table["Variable_Name"]
    )
    meta_data_table = meta_data_table[
        [
            "Variable_Name",
            "Variable_ID",
            "Table_ID",
            "Table_Name",
            "Type",
            "Unit",
            "Full_Name",
        ]
    ]

    # Sort by variable id
    meta_data_table.sort_values(by="Variable_ID", inplace=True)

    # Manually set Type to 'Count' for all tables
    meta_data_table["Type"] = "Count"

    # Save metadata table to CSV
    meta_data_table.to_csv(DATA_DIR / "Variable_Metadata.csv", index=False)


if __name__ == "__main__":
    download_scot_census_data()