from typing import Dict, Any, List, Optional
from datetime import datetime, timedelta
from pytz import timezone
from databricks.sdk import WorkspaceClient
from pyspark.sql import SparkSession
import sys
import time
import json
import os
import logging
import re
import pandas as pd
from .utils import (
    execute_run_databricks,
    execute_run_postgres,
    setup_logger,
    DbConnections
)

class LoadScripts:
    def __init__(self, base_dir: Optional[str] = None):
        self.logger = logging.setup_logger()

        self.config: Dict[str, Any] = {}
        self.template: Dict[str, Any] = {}
        self.rules: Dict[str, List[Dict[str, Any]]] = {}
        self.execution_id: Optional[int] = None
        self.job_start_time: Optional[float] = None
        
        self.databricks_client = None

        if base_dir:
            self.base_dir = base_dir

        elif os.getenv("BRV_BASE_DIR"):
            self.base_dir = os.getenv("BRV_BASE_DIR")
        else:
            self.base_dir = os.getcwd()
    
    def resolve_relative_path(self, raw_path: str) -> str:
        """
        Resolves a given relative or absolute file path using the base_dir.
        """
        raw_path = raw_path.strip().lstrip("/\\")
        return os.path.abspath(os.path.join(self.base_dir, raw_path))

    def load_config(self, config_path: str):
        with open(config_path, 'r') as f:
            self.config = json.load(f)

    def split_config_object(self):
        if 'RULES' in self.config:
            self.template = {k: v for k, v in self.config.items() if k != 'RULES'}
            self.rules = {'RULES': self.config['RULES']}
            return self.template, self.rules
        else:
            self.logger.warning("Array 'RULES' not found in rule object.")
            return self.template, None

    def run_exe(self, s: str) -> int:
        # Use consistent hashing and ensure non-negative number
        return abs(hash(s)) % (sys.maxsize + 1)

    def init_other_variables(self):
        # Use UTC date to avoid timezone issues
        hash_string = (datetime.utcnow() - timedelta(days=1)).strftime("%Y%m%d%H%M%S")
        self.execution_id = self.run_exe(hash_string)
        self.job_start_time = time.time()

    def set_system_variables(self, rule: Dict[str, Any]) -> Dict[str, Any]:
        self.init_other_variables()
        hash_string = f"{rule['TEAM_NAME']}{rule['DOMAIN_NAME']}{rule['RULE_CATEGORY_NAME']}{rule['RULE_ID']}"
        unique_rule_identifier = self.run_exe(hash_string)
        rule['EXECUTION_ID'] = self.execution_id
        rule['UNIQUE_RULE_IDENTIFIER'] = unique_rule_identifier
        rule['UNIQUE_TIMESTAMP'] = datetime.utcnow().strftime("%Y%m%d%H%M%S%f")
        return rule

    def replace_common_variables(self,sql_str, rule):
        """
        Replaces placeholders in sql_str with corresponding values from the rule dictionary.
        Any placeholder in the format <PLACEHOLDER_NAME> will be dynamically replaced.
        """
        final_sql = sql_str
        for key, value in rule.items():
            placeholder = f'<{key}>'
            final_sql = final_sql.replace(placeholder, str(value))
        return final_sql

    def get_file_text(self, file_path: str, is_local: bool = False) -> str:
        """Read file content from local filesystem or Databricks workspace.
        
        Args:
            file_path: Path to the file
            is_local: Whether the file is local or in Databricks workspace
            
        Returns:
            str: File contents
            
        Raises:
            FileNotFoundError: If local file not found
            Exception: For other errors
        """
        if is_local:
            try:
                # First, try to find SQL templates in the package directory
                package_dir = os.path.dirname(os.path.abspath(__file__))
                possible_paths = [
                    os.path.join(package_dir, 'jobs', file_path),  # Package sql_templates
                    os.path.join(os.getcwd(), 'jobs', file_path),  # Current directory sql_templates
                    os.path.join(self.base_dir, 'jobs', file_path),  # Base directory sql_templates
                ]

                # Add user-specified templates path if available
                if hasattr(self, 'jobs_path'):
                    possible_paths.insert(0, os.path.join(self.sql_templates_path, file_path))

                for path in possible_paths:
                    if os.path.exists(path):
                        full_path = path
                        break
                else:
                    raise FileNotFoundError(f"SQL template not found: {file_path}")

                self.logger.info(f"Reading local file: {full_path}")
                try:
                    with open(full_path, 'r', encoding='utf-8') as file:
                        return file.read()
                except UnicodeDecodeError:
                    self.logger.warning(f"UTF-8 decoding failed, retrying with latin-1 for file: {full_path}")
                    with open(full_path, 'r', encoding='latin-1') as file:
                        return file.read()

            except Exception as e:
                self.logger.error(f"Error reading file {file_path}: {str(e)}")
                raise
        else:
            # Download file content from Databricks workspace
            try:
                with self.databricks_client.workspace.download(file_path) as f:
                    return f.read().decode("utf-8")
            except Exception as e:
                self.logger.error(f"Error downloading from Databricks: {str(e)}")
                raise

    def set_business_rule_check_templates(self, rule: Dict[str, Any]) -> Dict[str, Any]:
        """Load SQL templates based on engine type.
        
        Args:
            rule: Rule configuration dictionary
            
        Returns:
            Updated rule with SQL templates
        """
        engine_type = rule.get('ENGINE_TYPE', '').lower()
        
        # First try package sql_templates directory
        package_dir = os.path.dirname(os.path.abspath(__file__))
        template_dir = os.path.join(package_dir, 'jobs')
        
        if engine_type == 'postgres':
            paths = {
                'batch_process': os.path.join('code_snippets', 'batch_process.sql'),
                'dynamic_variables': os.path.join('code_snippets', 'dynamic_variables.sql'),
                'generic_dwh_load_process': os.path.join('code_snippets', 'generic_dwh_load_process.sql'),
                'JOB_SCRIPTS': os.path.join('job_scripts', 'job_scripts.sql'),
            }
        else:
            raise ValueError(f"Unsupported ENGINE_TYPE in rule: '{engine_type}'")

        for key, path in paths.items():
            rule[key] = path

        snippet_keys = ['dynamic_variables', 'batch_process', 'generic_dwh_load_process']
        rule['CODE_SNIPPETS'] = [
            {'path': paths[key], 'order': idx}
            for idx, key in enumerate(snippet_keys, start=1)
        ]

        return rule


    def set_custom_template(self, rule: Dict[str, Any]) -> Dict[str, Any]:
        """Add SQL template paths to the rule based on engine type.
        
        Args:
            rule: Rule configuration dictionary
            
        Returns:
            Updated rule with SQL template paths and formatted SQL results
        """
        custom_template_path =  rule.get('JOB_SCRIPTS', '')
        if not custom_template_path:
            print("DEBUG: No JOB_SCRIPTS found in rule")
            self.logger.debug("No JOB_SCRIPTS found in rule, skipping adding SQL template")
            return rule

        try:
            print(f"DEBUG: Found JOB_SCRIPTS: {custom_template_path}")
            self.logger.info(f"Processing SQL template from: {custom_template_path}")
            
            # Remove leading slash if present and normalize path
            if custom_template_path.startswith('/'):
                custom_template_path = custom_template_path[1:]

            print(f"DEBUG: Normalized path: {custom_template_path}")

            # Read the SQL content from the file
            template_sql_content = self.get_file_text(custom_template_path, is_local=True)
            template_sql_content = template_sql_content.strip()

            # Store the SQL CONTENT, not the file path
            if 'CODE_SNIPPETS' in rule:
                print("DEBUG: Found CODE_SNIPPETS in rule, injecting snippets into template")
                self.logger.info("Injecting code snippets into SQL template")
                template_sql_content = self.inject_snippets_into_template(rule, template_sql_content)
                rule['JOB_SCRIPTS_SQL'] = template_sql_content
            else:
                print("DEBUG: No CODE_SNIPPETS found in rule, using raw template SQL")
                self.logger.info("Using raw SQL template without snippets")
                rule['JOB_SCRIPTS_SQL'] = template_sql_content

            print(f"DEBUG: SQL content loaded: {template_sql_content[:100]}...")  # First 100 chars
            self.logger.info(f"Executing SQL Template: {template_sql_content}")

            # Check if SQL content is empty
            if not template_sql_content:
                print("DEBUG: Template SQL content is empty")
                self.logger.warning("SQL template content is empty, skipping execution")
                return rule

        except FileNotFoundError as e:
            print(f"DEBUG: File not found error: {str(e)}")
            self.logger.error(f"SQL template file not found: {custom_template_path}. Error: {str(e)}")
        except Exception as e:
            print(f"DEBUG: General error: {str(e)}")
            self.logger.error(f"Error processing SQL template: {str(e)}")    
        
        return rule


    def extract_snippet_blocks(self, sql_text: str) -> Dict[str, str]:
        """
        Extracts SQL code blocks from a snippet file using markers like:
            #START <TAG>
            ...SQL code...
            #END <TAG>
        Returns a dict: { "TAG": "…sql code…" }
        """
        pattern = r"#START\s+(#EXEC\s+)?<(?P<tag>[^>]+)>\s*(?P<code>.*?)#END\s+<\2>"
        #pattern = r"#START\s+<(?P<tag>[^>]+)>\s*(?P<code>.*?)#END\s+<\1>"
        matches = re.finditer(pattern, sql_text, re.DOTALL | re.IGNORECASE)
        blocks = {}
        for match in matches:
            exec_flag = bool(match.group(1))
            tag = match.group("tag").strip()
            code = match.group("code").strip()
            blocks[tag] = {"code": code, "exec": exec_flag}
        return blocks

    def execute_sql_scalar(self, sql_text: str, engine_type: str) -> Any:
        """
        Runs the given SQL as a scalar query and returns the first cell.
        Handles Databricks or Postgres.
        """
        engine_type = engine_type.lower()

        if engine_type == 'postgres':
            # Postgres: use your helper
            db_connection = DbConnections()
            engine = db_connection.postgres_engine()
            result = execute_run_postgres(sql_text, engine, retrieve_result=True)
            if result and isinstance(result, list) and len(result) > 0:
                rows = result[0][1:]  # skip header
                if rows:
                    return rows[0][0]
            return None



    def assign_snippet_blocks_to_rule(self, rule: Dict[str, Any]) -> Dict[str, Any]:
        """
        Assigns values from SQL snippet blocks into the rule dictionary.
        For example:
            #START <SCHEMA>
            seacomp
            #END <SCHEMA>
        becomes rule["SCHEMA"] = "seacomp"

        Args:
            rule: Rule dictionary containing CODE_SNIPPETS.

        Returns:
            Updated rule with direct assignments from snippet blocks.
        """
        snippet_blocks = {}
        engine_type = rule.get("ENGINE_TYPE", "postgres")

        # Safely get and sort CODE_SNIPPETS
        snippets = sorted(rule.get("CODE_SNIPPETS", []), key=lambda x: x.get("order", 0))

        for snippet_info in snippets:
            path = snippet_info.get("path")
            if not path:
                continue
            try:
                resolved_path = self.resolve_relative_path(path)
                with open(resolved_path, 'r') as f:
                    sql_text = f.read()
                    blocks = self.extract_snippet_blocks(sql_text)
                    for tag, block_info in blocks.items():
                        if block_info["exec"]:
                            sql_to_run = self.replace_common_variables(block_info["code"], rule)
                            result = self.execute_sql_scalar(sql_to_run, engine_type)
                            snippet_blocks[tag] = result.strip() if isinstance(result, str) else (str(result).strip() if result else "")
                        else:
                            snippet_blocks[tag] = block_info["code"] 
                   
            except FileNotFoundError:
                print(f"DEBUG: Snippet file not found: {path} -- Skipping.")
                self.logger.warning(f"Snippet file not found: {path} -- Skipping.")

        # Update rule with block values
        rule.update(snippet_blocks)
        return rule


    def inject_snippets_into_template(self, rule: Dict[str, Any], template_sql: str) -> str:
        """
        Replaces placeholders in template_sql with snippets defined in rule["CODE_SNIPPETS"].
        Placeholders are:
        - #START <TAG>...#END <TAG>
        - or <TAG>
        """
        snippet_blocks = {}

        # Sort snippets by 'order'
        snippets = sorted(rule.get("CODE_SNIPPETS", []), key=lambda x: x["order"])

        for snippet_info in snippets:
            path = snippet_info["path"]
            try:
                resolved_path = self.resolve_relative_path(path)
                with open(resolved_path, 'r') as f:
                    sql_text = f.read()
                    blocks = self.extract_snippet_blocks(sql_text)
                    snippet_blocks.update(blocks)
            except FileNotFoundError as e:
                print(f"Snippet file not found: {path} -- Skipping.")

        # Replace placeholders in template_sql
        for tag, block_info in snippet_blocks.items():
            code = block_info["code"] 

            # Replace full block if exists
            block_pattern = rf"#START <{tag}>.*?#END <{tag}>"
            if re.search(block_pattern, template_sql, re.DOTALL):
                template_sql = re.sub(block_pattern, code, template_sql, flags=re.DOTALL)
            else:
                # Replace tag placeholder like <TAG>
                template_sql = template_sql.replace(f"<{tag}>", code)

        return template_sql
    
    def set_custom_template_queries(self, rule: dict) -> dict:
        rule['JOB_SCRIPTS_SQL'] = self.replace_common_variables(rule.get('JOB_SCRIPTS_SQL', ''), rule)

        return rule

    
    def exec_custom_template_check_final_queries(self, rule):
        engine_type = rule.get("ENGINE_TYPE", "").lower()
        if engine_type == 'databricks':
            print("Running Databricks...")
            try:
                execute_run_databricks(rule["JOB_SCRIPTS_SQL"], False)
            except Exception as e:
                print(f"Error executing JOB_SCRIPTS_SQL on Databricks: {str(e)}")

        else:  # Assume Postgres
            print("Running Postgres...")
            with open(f"generated_{rule['RULE_NAME']}.sql", "w") as f:
                f.write(rule["JOB_SCRIPTS_SQL"])
            db_connection = DbConnections()
            engine = db_connection.postgres_engine()
            try:
                execute_run_postgres(rule["JOB_SCRIPTS_SQL"], engine, False,load_script=True)
            except Exception as e:
                print(f"Error executing JOB_SCRIPTS_SQL on Postgres: {str(e)}")

        return rule
    
    def business_rule_check(self) -> Dict[str, List[Dict[str, Any]]]:

        if self.template.get("RULE_CATEGORY_NAME") == 'JOB_SCRIPTS':
            for idx, rule in enumerate(self.rules.get("RULES", [])):
                rule = {**self.template, **rule}
                rule = self.set_system_variables(rule)
                rule = self.assign_snippet_blocks_to_rule(rule)
                rule = self.set_custom_template(rule)
                rule = self.set_custom_template_queries(rule)
                rule = self.exec_custom_template_check_final_queries(rule)
                
                self.rules["RULES"][idx] = rule
                self.logger.info(json.dumps(rule, indent=2))

        return self.rules
