from flask import Blueprint, request, jsonify
import logging
import uuid
import pytz
import json

import pandas as pd
import numpy as np
from google.cloud import bigquery
from firebase.firebase import Firebase
from collections import Counter
from datetime import datetime
from bigquery.bq import BigQuery, Schema
from utility.function import Function, Pack, Key, Profile
from utility.token import Token
from utility.authorization import TokenAuthorization
logging.basicConfig(level=logging.INFO)

schema = Blueprint('schema', __name__, url_prefix='/schema')
fb = Firebase(host="https://ydm-profile-360-default-rtdb.asia-southeast1.firebasedatabase.app/")
timezone = pytz.timezone('Asia/Bangkok')

@schema.route('/mapping', methods=['GET','POST'])
def mapping():
    if request.method == 'GET':
        offline_mapping = fb.db.reference(f'template/offline_mapping').get()
        mapping_option = list(set(item for sublist in offline_mapping.values() for item in sublist))
        return {"status":"ok",'message': 'Success', 'data': mapping_option}, 200
    
    return {"status":"error",'message': "Method not allowed"}, 405

@schema.route('/list', methods=['GET'])
def list_schema():
    local_key = request.headers.get('Authorization')
    bearer = local_key.replace("Bearer ", "")
    userAuth = TokenAuthorization(access_token=bearer)
    userData = userAuth.validate()
    
    if request.method == 'GET':
        
        param = request.args
        for req in ['property_id']:
            if req not in param:
                return {"status":"error",'message': f"{req} is require"}, 400
            
        property_id = request.args.get('property_id')
        
        offline_schema = fb.db.reference(f'property/{property_id}/offline_mapping').get()
        summaries = []
        
        for name, cfg in offline_schema.items():
            if not cfg:  # skip empty dicts
                continue
            summaries.append({
                "name": name,
                "row_count": cfg.get("row_count", 0),
                "status": cfg.get("status"),
                "createdate": cfg.get("createdate"),
                "lastupdate": cfg.get("lastupdate"),
            })
        return {"status":"ok",'message': 'Success', 'data': summaries}, 200
    return {"status":"error",'message': "Method not allowed"}, 405

@schema.route('/get_schema', methods=['GET'])
def get_schema():
    if request.method == 'GET':
        
        param = request.args
        for req in ['property_id', 'table_name']:
            if req not in param:
                return {"status":"error",'message': f"{req} is require"}, 400
            
        property_id = request.args.get('property_id')
        table_name = request.args.get('table_name')
        
        bq = BigQuery()
        schema = bq.get_offline_table_schema_as_bqjson(property_id,table_name)
        return {"status":"ok",'message': 'Success', 'data': schema}, 200
    
    return {"status":"error",'message': "Method not allowed"}, 405


@schema.route('/get_value', methods=['GET'])
def get_value():
    param = request.args
    for req in ["property_id", "table_name", "field_name"]:
        if req not in param or not str(param[req]).strip():
            return jsonify({"status":"error","message": f"{req} is require"}), 400

    bq = BigQuery()
    result = bq.get_field_values(
        property_id=param["property_id"].strip(),
        table_name=param["table_name"].strip(),   # usually just the table name
        field_path=param["field_name"].strip(),
        distinct_limit=int(param.get("limit", 1000))
    )
    code = 400 if result.get("status") == "error" else 200
    return jsonify(result), code

        
@schema.route('/manage', methods=['GET','POST','PUT','DELETE'])
def manage_schema():
    local_key = request.headers.get('Authorization')
    bearer = local_key.replace("Bearer ", "")
    userAuth = TokenAuthorization(access_token=bearer)
    userData = userAuth.validate()
    
    if request.method == 'GET':
        
        param = request.args
        for req in ['property_id', 'table_name']:
            if req not in param:
                return {"status":"error",'message': f"{req} is require"}, 400
            
        property_id = request.args.get('property_id')
        table_name = request.args.get('table_name')
        
        bq = BigQuery()
        schema = bq.get_schema_dataset(property_id,table_name)
        offline_mapping = fb.db.reference(f'property/{property_id}/offline_mapping/offline_{table_name}').get()
        
        if offline_mapping:
            schema_mapping = Schema.augment_schema_with_mapping_and_unique(schema, offline_mapping)
            return {"status":"ok",'message': 'Success', 'data': schema_mapping}, 200
        
        return {"status":"error",'message': f'Table {table_name} not found'}, 200
    
    elif request.method == 'POST':
        data = request.get_json()
        
        for req in ['property_id','table_name','schema']:
            if req not in data:
                return {"status":"error",'message': f"{req} is require for creating schema"}, 400
            
        offline_mapping = fb.db.reference(f'template/offline_mapping').get()
        table_name = data.get('table_name', None)
        property_id = data.get('property_id', None)
        schema = data.get('schema', [])
        
        if len(schema) < 2:
            return {"status":"error",'message': "Atleast 2 columns is required to create schema"}, 400
        
        table_exist = fb.db.reference(f'property/{property_id}/offline_mapping/offline_{table_name}').get()
        if table_exist:
            return {"status":"error",'message': f"Table {table_name} already exists"}, 400
        
        # process schema by mapping key from template
        result = Schema.process_schema(schema, offline_mapping)

        schema = result.get('schema', None)
        partition = result.get('partition',None)
        unique_col = result.get('unique_col', None)
        contact_key = result.get("contact_key", None)
        demographic_col = result.get("demographic_col", None)
        sensitive_data = result.get("sensitive_data", None)
        other_key = result.get("other_key", None)
        property_id = data.get('property_id', None)
        table_id = data.get('table_name', None)

        schema.append({"name": "version_key", "type": "DATETIME", "mode": "REQUIRED","description":"The version key of the data, used to track changes in the schema or data structure."})
        schema.append({"name": "user_pseudo_id", "type": "STRING", "mode": "NULLABLE", "description":"A unique identifier for each user, used to track user behavior and interactions."})
        schema.append({"name": "ref_user_pseudo_id", "type": "STRING", "mode": "NULLABLE", "description":"(REF) A unique identifier for each user, used to track user behavior and interactions."})
        
        bq = BigQuery()
        bq.create_table_from_json_schema(f"client_{property_id}",f"offline_{table_id}",schema,partition)
        bq.create_table_from_json_schema(f"client_{property_id}",f"offline_{table_id}_temp",schema,partition)
        
        offline_mapping = {}

        if contact_key:
            offline_mapping['contact_key'] = contact_key
        if other_key:
            offline_mapping['other_key'] = other_key
        if partition:
            offline_mapping['date_column'] = partition
        if demographic_col:
            offline_mapping['demographic_col'] = demographic_col
        if sensitive_data:
            offline_mapping['sensitive_data'] = sensitive_data
        if unique_col:
            offline_mapping['unique_col'] = unique_col

        timezone = pytz.timezone('Asia/Bangkok')
        offline_mapping['createdate'] = datetime.now(timezone).strftime('%Y-%m-%d %H:%M:%S')
        offline_mapping['lastupdate'] = datetime.now(timezone).strftime('%Y-%m-%d %H:%M:%S')
        offline_mapping['row_count'] = 0
        offline_mapping['status'] = 'active'
        
        fb.db.reference(f'property/{property_id}/offline_mapping/offline_{table_id}').set(offline_mapping)
        
        return {"status":"Create schema successfully"}, 200
    
    elif request.method == 'DELETE':
        
        param = request.args
        
        for req in ['property_id','table_id']:
            if req not in param:
                return {"status":"error",'message': f"{req} is require"}, 400
        
        property_id = request.args.get('property_id')
        table_id = request.args.get('table_id')
        
        bq = BigQuery()
        bq.delete_table("customer-360-profile",f"client_{property_id}",f"offline_{table_id}")
        bq.delete_table("customer-360-profile",f"client_{property_id}",f"offline_{table_id}_temp")
        
        ## delete offline mapping from property
        fb.db.reference(f'property/{property_id}/offline_mapping/offline_{table_id}').delete()
        return {"status": "success", "message": "Table and mapping deleted"}, 200  # <-- ADD THIS
    elif request.method == 'PUT':
        data = request.get_json()
        
        for req in ['property_id','table_name','schema']:
            if req not in data:
                return {"status":"error",'message': f"{req} is require for creating schema"}, 400
        
        offline_mapping = fb.db.reference(f'template/offline_mapping').get()
        schema = data.get('schema', [])
        
        if len(schema) < 2:
            return {"status":"error",'message': "Atleast 2 columns is required to create schema"}, 400
        result = Schema.process_schema(schema, offline_mapping)

        schema = result.get('schema', None)
        partition = result.get('partition',None)
        unique_col = result.get('unique_col', None)
        contact_key = result.get("contact_key", None)
        demographic_col = result.get("demographic_col", None)
        sensitive_data = result.get("sensitive_data", None)
        other_key = result.get("internal_key", None)

        property_id = data.get('property_id', None)
        table_id = data.get('table_name', None)

        schema.append({"name": "version_key", "type": "STRING", "mode": "REQUIRED","description":"The version key of the data, used to track changes in the schema or data structure."})
        
        bq = BigQuery()
        bq.delete_table("customer-360-profile",f"client_{property_id}",f"offline_{table_id}")
        bq.create_table_from_json_schema(f"client_{property_id}",f"offline_{table_id}",schema,partition)
        
        createDate = fb.db.reference(f'property/{property_id}/offline_mapping/offline_{table_id}/createdate').get()
        
        offline_mapping = {}

        if contact_key:
            offline_mapping['contact_key'] = contact_key
        if other_key:
            offline_mapping['other_key'] = other_key
        if partition:
            offline_mapping['date_column'] = partition
        if demographic_col:
            offline_mapping['demographic_col'] = demographic_col
        if sensitive_data:
            offline_mapping['sensitive_data'] = sensitive_data
        if unique_col:
            offline_mapping['unique_col'] = unique_col

        
        timezone = pytz.timezone('Asia/Bangkok')
        offline_mapping['createdate'] = createDate
        offline_mapping['lastupdate'] = datetime.now(timezone).strftime('%Y-%m-%d %H:%M:%S')
        offline_mapping['row_count'] = 0
        offline_mapping['status'] = 'active'
        
        fb.db.reference(f'property/{property_id}/offline_mapping/offline_{table_id}').set(offline_mapping)
        
        return {"status":"Edit schema successfully"}, 200
    return {"status":"error",'message': "Method not allowed"}, 405

@schema.route('/column', methods=['GET'])
def get_column_schema():
    if request.method == 'GET':
        
        param = request.args
        for req in ['property_id', 'table_name']:
            if req not in param:
                return {"status":"error",'message': f"{req} is require"}, 400
            
        property_id = request.args.get('property_id')
        table_name = request.args.get('table_name')
        
        bq = BigQuery()
        columns = bq.get_bq_columns(property_id, table_name)
        
        if len(columns) == 0:
            return {"status":"error",'message': f'Table {table_name} not found'}, 404
        
        return {"status":"ok",'message': 'Success', 'data': columns}, 200
    
    return {"status":"error",'message': "Method not allowed"}, 405

@schema.route('/export', methods=['GET'])
def export_schema():
    if request.method == 'GET':
        
        param = request.args
        for req in ['property_id', 'table_name']:
            if req not in param:
                return {"status":"error",'message': f"{req} is require"}, 400
            
        property_id = request.args.get('property_id')
        table_name = request.args.get('table_name')
        
        bq = BigQuery()
        schema = bq.get_schema_dataset(property_id,table_name)
        return {"status":"ok",'message': 'Success', 'data': schema}, 200
        
    return {"status":"error",'message': "Method not allowed"}, 405

@schema.route('/edit/status', methods=['PUT'])
def edit_status():
    if request.method == 'PUT':
        data = request.get_json()
        
        for req in ['property_id','table_name','status']:
            if req not in data:
                return {"status":"error",'message': f"{req} is require for detecting schema"}, 400
        
        status = data.get('status', None)
        property_id = data.get('property_id', None)
        table_name = data.get('table_name', None)
        
        current_status = fb.db.reference(f'property/{property_id}/offline_mapping/offline_{table_name}/status').get()   
        if status == current_status:
            return {"status":"error",'message': f"Status already set to {status}"}, 400
        else:
            timezone = pytz.timezone('Asia/Bangkok')
            fb.db.reference(f'property/{property_id}/offline_mapping/offline_{table_name}/status').set(status)
            fb.db.reference(f'property/{property_id}/offline_mapping/offline_{table_name}/lastupdate').set(datetime.now(timezone).strftime('%Y-%m-%d %H:%M:%S'))
            return {"status":"ok",'message': f"Schema {table_name}'s status updated to {status}"}, 200
        
    return {"status":"error",'message': "Method not allowed"}, 405

        
@schema.route('/sample/data', methods=['POST'])
def preview_data():
    if request.method == 'POST':
        data = request.get_json()
        
        for req in ['example_data']:
            if req not in data:
                return {"status":"error",'message': f"{req} is require for detecting schema"}, 400
        
        example_data = data["example_data"]
        sample = example_data[:5] if isinstance(example_data, list) else []
        return {"status": "ok", "sample": sample}, 200

    return {"status":"error",'message': "Method not allowed"}, 405


@schema.route('/detect', methods=['GET','POST'])
def detect_schema():
    if request.method == 'POST':
        data = request.get_json()
        
        for req in ['example_data']:
            if req not in data:
                return {"status":"error",'message': f"{req} is require for detecting schema"}, 400
            
        example_data = data.get('example_data', None)
        schema = Schema.infer_bq_schema_lite(example_data)
        
        return {"status": "ok", "message": "Schema detected successfully", "data": schema}, 200
    return {"status":"error",'message': "Method not allowed"}, 405

@schema.route('/rename_column', methods=['GET','POST','DELETE'])
def rename_column():
    if request.method == 'POST':
        data = request.get_json()
        
        timezone = pytz.timezone('Asia/Bangkok')
        for req in ['new_name','old_name','property_id','table_id']:
            if req not in data:
                return {"status":"error",'message': f"{req} is require for creating schema"}, 400
            
        new_name = data.get('new_name',"")
        old_name = data.get('old_name',"")
        property_id = data.get('property_id','')
        table_id = data.get('table_id','')
        
        query = f"""
        ALTER TABLE `client_{property_id}.offline_{table_id}`
        RENAME COLUMN {old_name} TO {new_name}
        """
        try:
            bq = BigQuery()
            query_job = bq.client.query(query)
            query_job.result()
        
            ref = fb.db.reference(f'property/{property_id}/offline_mapping/offline_{table_id}')
            mapping_data = ref.get()

            for group in ["contact_key", "other_key"]:
                group_map = mapping_data.get(group, {})
                for schema_key, mapped_value in group_map.items():
                    if mapped_value == old_name:
                        # Perform the update: use the original schema_key
                        ref.child(group).child(schema_key).set(new_name)
                        break
            
            ref.update({"lastupdate": datetime.now(timezone).strftime("%Y-%m-%d %H:%M:%S")})
            return {"status":"Edit schema successfully"}, 200
        
        except:
            return {"status":"error",'message': "Old Name doesn't match with existing field."}, 100
        
    return {"status":"error",'message': "Method not allowed"}, 405


@schema.route('/add_column', methods=['GET','POST','DELETE'])
def add_column():
    if request.method == 'POST':
        data = request.get_json()
        
        for req in ['property_id','table_id','schema']:
            if req not in data:
                return {"status":"error",'message': f"{req} is require for creating schema"}, 400
            
        property_id = data.get("property_id",None)
        table_id = data.get("table_name", None)
        schema = data.get("schema", None)
        
        table_id = f"customer-360-profile.client_{property_id}.offline_{table_id}"
        new_schema_fields = [bigquery.SchemaField(**field) for field in schema]
        
        bq = BigQuery()
        table = bq.client.get_table(table_id)
        updated_schema = table.schema + new_schema_fields
        table.schema = updated_schema
        bq.client.update_table(table, ["schema"])
        
        return {"status": f"Add Column to {table_id} successfully"}, 200
        
    return {"status":"error",'message': "Method not allowed"}, 405


@schema.route('/delete_column', methods=['GET','POST','DELETE'])
def delete_column():
    if request.method == 'POST':
        data = request.get_json()
        
        for req in ['property_id','table_id',"drop_columns"]:
            if req not in data:
                return {"status":"error",'message': f"{req} is require for creating schema"}, 400
            
        property_id = data.get("property_id",None)
        table_id = data.get("table_name", None)
        drop_columns = data.get("drop_columns",[])
        
        if len(drop_columns) == 0:
            return {"status":"error",'message': "Column name required for drop_columns"}, 400
            
        table_id = f"customer-360-profile.client_{property_id}.offline_{table_id}"
        
        bq = BigQuery()
        for col in drop_columns:
            query = f"""
            ALTER TABLE `{table_id}`
            DROP COLUMN {col}
            """
            bq.client.query(query).result()
        
        return {"status": f"Drop Column {drop_columns} from {table_id} successfully"}, 200
        
    return {"status":"error",'message': "Method not allowed"}, 405