import os
import logging
import datetime
import json
import ast
import jwt
import urllib #Needed for avoiding using requests, which modified the S3 presigned URL
from urllib.parse import urlparse, parse_qs
import csv
import pathlib

from .utils import convert_timestamps, send_post_request, send_get_request, decode_jwt_token, querystr_to_json, att2statement
from .lineage_graph import create_lineage_graph 

from .constants import (
    SCRIBE_VULNERABILITIES_DATASET,
    SCRIBE_PRODUCTS_DATASET,
    SCRIBE_POLICY_DATASET,
    SCRIBE_TEAM_STAT_DATASET,
    SCRIBE_LINEAGE_DATASET,
    SCRIBE_RISK_DATASET,
    SCRIBE_FINDINGS_DATASET
)

class ScribeClient:
    scribe_vulnerabilities_dataset = SCRIBE_VULNERABILITIES_DATASET
    scribe_products_dataset = SCRIBE_PRODUCTS_DATASET
    scribe_policy_dataset = SCRIBE_POLICY_DATASET
    scribe_team_stat_dataset = SCRIBE_TEAM_STAT_DATASET
    scribe_lineage_dataset=SCRIBE_LINEAGE_DATASET
    scribe_risk_dataset = SCRIBE_RISK_DATASET
    scribe_findings_dataset = SCRIBE_FINDINGS_DATASET

    def get_chache_data(self):
        # Exclude product_list from cache as per requirements
        return {
            "timestamp": datetime.datetime.now().isoformat(),
            "datasets": self.datasets,
            "team_name": self.team_name,
            "team_info": self.team_info
        }
    
    def need_cache_data_update(self, cached_data):
        # if timestamp older than one day
        if cached_data and "timestamp" in cached_data:
            try:
                last_update = datetime.datetime.fromisoformat(cached_data["timestamp"])
            except Exception as e:
                logging.error(f"Error parsing timestamp: {e}")
                return True
            return (datetime.datetime.now() - last_update).days > 1
        
    def __init__(self, api_token=None, ignore_env = False, base_url="https://api.scribesecurity.com", refresh_interval=600, cached_data=None, env = "prod",
             username=None, password=None):
    

        def update_from_cached_data(cached_data):
            if "datasets" in cached_data:
                self.datasets = cached_data["datasets"]
                self.dataset_ids = self.create_dataset_id_map(self.datasets)
            else:
                logging.error("No datasets in provided cached data")
            # product_list intentionally not loaded from cache
            if "team_name" in cached_data:
                self.team_name = cached_data["team_name"]
            else:
                logging.error("No team name in provided cached data")
            if "team_info" in cached_data:
                self.team_info = cached_data["team_info"]
            else:
                logging.error("No team info in provided cached data")

        # Default cache file path
        cache_dir = os.path.join(os.getcwd(), ".scribe")
        cache_file = os.path.join(cache_dir, "cache.json")
        loaded_cache = None

        # Only load cache if not explicitly provided
        if cached_data is None:
            try:
                if os.path.exists(cache_file):
                    # Check file age
                    file_mtime = datetime.datetime.fromtimestamp(os.path.getmtime(cache_file))
                    if (datetime.datetime.now() - file_mtime).days <= 1:
                        with open(cache_file, "r") as f:
                            loaded_cache = json.load(f)
                    else:
                        logging.info("Cache file is older than one day, will refresh cache.")
                else:
                    # Ensure .scribe directory exists
                    pathlib.Path(cache_dir).mkdir(parents=True, exist_ok=True)
            except Exception as e:
                logging.error(f"Error loading cache file: {e}")
        else:
            loaded_cache = cached_data

        if ignore_env:
            self.api_token = api_token
        else:
            self.api_token = api_token or os.getenv('SCRIBE_TOKEN', None)
        if not self.api_token:
            logging.error("No API token provided. Please set the SCRIBE_TOKEN environment variable or pass it as an argument.")
            return
        self.username = username
        self.password = password
        self.env = env
        self.base_url = base_url
        
        if env == "prod":
            self.superset_base_url = "https://superset.scribesecurity.com"
        else:
            self.superset_base_url = f"https://internal-superset.{env}.scribesecurity.com"
        
        self.refresh_interval = refresh_interval
        self.last_refresh = datetime.datetime.now() - datetime.timedelta(days=5)  # Initialize to a time in the past

        self.jwt_token = None
        self.superset_token = None
        self.dataset_ids = None
        self.datasets = None
        self.product_list = None
        self.team_id = None
        self.team_name = None
        self.team_info = None
        p = None
        if self.api_token:
            self.last_refresh = datetime.datetime.now() # - datetime.timedelta(hours=5)
            self.refresh_data(force=True)
            
            # Use loaded_cache if available and not stale, else refresh and save
            if not loaded_cache or self.need_cache_data_update(loaded_cache):
                d = self.get_datasets()
                if not d:
                    logging.error("Refresh: Failed to update dataset IDs")
                else:
                    self.datasets = d
                    self.dataset_ids = self.create_dataset_id_map(d)
                    p = self.get_products(force=True)
                if not p:
                    logging.error("Refresh: Failed to update product list")
                    logging.error(f"Using previous product list\n{self.product_list}")
                else:
                    self.product_list = p

                t = self.get_team_info()
                if not t:
                    logging.error("Failed to get team data.")

                # Save cache (excluding product_list)
                try:
                    pathlib.Path(cache_dir).mkdir(parents=True, exist_ok=True)
                    cache_data = self.get_chache_data()
                    with open(cache_file, "w") as f:
                        json.dump(cache_data, f, indent=2)
                except Exception as e:
                    logging.error(f"Error saving cache file: {e}")
            else:
                update_from_cached_data(loaded_cache)

    def validate_dataset(self, dataset_name):
        if not self.dataset_ids:
            self.refresh_data()
        if not self.dataset_ids:
            logging.error("Failed to get dataset IDs")
            return False
        if dataset_name not in self.dataset_ids:
            logging.error(f"Dataset {dataset_name} not found")
            return False
        return True
    def set_api_token(self, api_token):
        if self.is_api_token_set():
            return
        self.api_token = api_token
        if self.api_token:
            self.last_refresh = datetime.datetime.now()
            self.refresh_data(force=True)

    def is_api_token_set(self):
        return self.api_token is not None


    def refresh_data(self, force=False):
        if not self.api_token:
            logging.error("No API token provided")
            return
        
        need_refresh = (
            force or
            not self.dataset_ids or
            not self.product_list or
            (datetime.datetime.now() - self.last_refresh).seconds > self.refresh_interval
        )

        if not need_refresh:
            return

        self.last_refresh = datetime.datetime.now()

        self.jwt_token = self.login(self.api_token, self.base_url)
        if not self.jwt_token:
            logging.error("Refresh: Failed to get JWT token")

        try:
            decoded_token = jwt.decode(self.jwt_token, options={"verify_signature": False})
            self.decoded_token = decoded_token
        except jwt.InvalidTokenError as e:
            logging.error(f"Invalid JWT token: {e}")
            self.decoded_token = None
            return
        jwt_name = decoded_token.get("name")
        try:
            prefix = "scribe-hub-team" if self.env == "prod" else f"scribe-hub-{self.env}-team"
            self.team_id = jwt_name.split(prefix)[1].split("-")[0]
            logging.info(f"Team:{self.team_id}")
        except Exception as e:
            logging.error(f"Error extracting team ID: {e}")
            self.team_id = None


        self.superset_token = self.get_superset_token(self.jwt_token, self.base_url)
        if not self.superset_token:
            logging.error("Refresh: Failed to get Superset token")

    def send_post_request(self, url, token=None, body=None):
        """
        Deprecated: Use `send_post_request` from `utils.py` instead.
        """
        return send_post_request(url, token, body)

    def send_get_request(self, url, token=None):
        """
        Deprecated: Use `send_get_request` from `utils.py` instead.
        """
        return send_get_request(url, token)

    def decode_jwt_token(self, jwt_token):
        """
        Deprecated: Use `decode_jwt_token` from `utils.py` instead.
        """
        return decode_jwt_token(jwt_token)

    def querystr_to_json(self, query):
        """
        Deprecated: Use `querystr_to_json` from `utils.py` instead.
        """
        return querystr_to_json(query)

    def att2statement(self, data):
        """
        Deprecated: Use `att2statement` from `utils.py` instead.
        """
        return att2statement(data)

    def login(self, api_token, base_url):
        url = f"{base_url}/v1/login"
        body = {"api_token": api_token}
        response = self.send_post_request(url=url, body=body)

        if response and 'token' in response.json():
            return response.json()['token']
        else:
            logging.error("Token not found in the response.")
            return None

    def get_superset_token(self, jwt_token, base_url):
        url = f"{base_url}/dataset/token"
        response = self.send_get_request(url=url, token=jwt_token)
        if not response:
            return None

        response_json = json.loads(response.text)
        if 'access_token' in response_json:
            return response_json['access_token']
        else:
            logging.error("Access token not found in the response.")
            return None

    def create_dataset_id_map(self, datasets):
        data = datasets
        if not data:
            return {}

        try:
            result = data.get("result", [])
            datasource_id_map = {}
            for entry in result:
                datasource_name = entry.get("datasource_name")
                datasource_id = entry.get("id")
                if datasource_name and datasource_id:
                    datasource_id_map[datasource_name] = entry
            return datasource_id_map
        except json.JSONDecodeError as e:
            logging.error(f"Error decoding JSON: {e}")
            return {}

    def get_team_info(self):
        if not self.team_id:
            logging.error("Team ID not set.")
            return False
        if not self.validate_dataset(self.scribe_team_stat_dataset):
            return False
        query = {
            "columns": [
                "teamName"
            ],
            "filters": [
                {
                    "col": "teamId",
                    "op": "==",
                    "val": self.team_id
                }
            ],
            "metrics": [],
            "orderby": [],
            "row_limit": 1
        }
        r = self.query_superset(self.scribe_team_stat_dataset, query)
        try:
            self.team_name = r["result"][0]["data"][0]["teamName"]
            self.team_info = r["result"][0]["data"][0]
        except (KeyError, IndexError) as e:
            logging.error(f"Error accessing response data: {e}")
            return False
        except json.JSONDecodeError as e:
            logging.error(f"Error decoding JSON: {e}")
            return False

        return True
    
    def get_link(self, use = "products", text = "Link", params = {}):
        params_str = ""
        if params:
            # %3B is the URL encoded version of ';'
            params_str = ";;".join([f"{k}:{v}" for k, v in params.items()])
            params_str = "searchFilters=" + params_str

        if use == "products":
            return f"\n[{text}](https://scribe-security.github.io/redirect/redirect.html?path=%2Fproducer-products)"
        # Add here elif for all uses
        elif use == "product_vulnerabilities":
            if params:
                params_str = '&' + params_str
            # return f"\n[{text}](http://localhost:8090/redirect.html?path=%2Fsbom?redirectTabName=Vulnerabilities&{params_str})"
            return f"\n[{text}](https://scribe-security.github.io/redirect/redirect.html?path=%2Fsbom?redirectTabName=Vulnerabilities{params_str})"
        elif use == "product_policy_results":
            if params:
                params_str = '?' + params_str
            return f"\n[{text}](https://scribe-security.github.io/redirect/redirect.html?path=%2Fpolicy%2Fevaluation{params_str})"
        else:
            return f"\n[{text}](https://app.scribesecurity.com)"

    def get_products(self, all_versions=False, force=False):
        if not force:
            if (datetime.datetime.now() - self.last_refresh).seconds < self.refresh_interval and self.product_list:
                return self.product_list
        if not self.validate_dataset(self.scribe_products_dataset):
            return None

        url = f"{self.base_url}/dataset/data"
        body = {
            "superset_token": self.superset_token,
            "validate": False,
            "query": {
                "datasource": {
                    "id": self.dataset_ids[self.scribe_products_dataset]["id"],
                    "type": "table"
                },
                "force": "false",
                "queries": [
                    {
                        "columns": [
                            "logical_app",
                            "logical_app_version",
                            "version_timestamp"
                        ],
                        "filters": [],
                        "metrics": [],
                        "orderby": [["logical_app", True], ["version_timestamp", False]],
                        "row_limit": 0
                    }
                ],
                "result_format": "json",
                "result_type": "results"
            }
        }

        r = self.send_post_request(url=url, token=self.jwt_token, body=body)
        if not r:
            return self.product_list

        try:
            o = r.json()
            o = convert_timestamps(o, "version_timestamp")
            result = []
            for product in o["result"][0]["data"]:
                item = {"name": product["logical_app"], "version": product["logical_app_version"], "timestamp": product["version_timestamp"]}
                result.append(item)
            self.product_list = result
        except Exception as e:
            logging.error(f"Error decoding JSON: {e}")
            return self.product_list

        if all_versions:
            return result

        # Keep only latest version per product
        latest = []
        last_product = ""
        for product in result:
            if product['name'] == last_product:
                continue
            last_product = product['name']
            latest.append(product)

        # sort by timestamp
        latest.sort(key=lambda x: x['timestamp'], reverse=True)
 
        return latest

    def format_table(self, headers, rows):
        """
        Format a table with aligned columns for better readability in plain text.

        Parameters:
            headers (list): List of column headers.
            rows (list of lists): List of rows, where each row is a list of column values.

        Returns:
            str: Formatted table as a string.
        """
        # Calculate the maximum width of each column
        column_widths = [len(header) for header in headers]
        for row in rows:
            for i, value in enumerate(row):
                column_widths[i] = max(column_widths[i], len(str(value)))

        # Create a format string for each row
        row_format = " | ".join(f"{{:<{width}}}" for width in column_widths)

        # Format the header and rows
        header_line = row_format.format(*headers)
        separator_line = "-+-".join("-" * width for width in column_widths)
        row_lines = [row_format.format(*[str(value) for value in row]) for row in rows]

        # Combine all parts into a single string
        return "\n".join([header_line, separator_line] + row_lines)

    def format_query_result_table(self, query_result):
        """
        Format query results into an aligned table for better readability.

        Parameters:
            query_result (dict): The query result containing columns and data.

        Returns:
            str: Formatted table as a string.
        """
        try:
            data = query_result["result"][0]["data"]
            if not data:
                return "No data found"

            columns = query_result["result"][0]["colnames"]
            if not columns:
                return "No columns information found"

            # Prepare headers and rows for the table
            headers = columns
            rows = [[row[col] for col in columns] for row in data]

            # Format the table using the existing format_table method
            return self.format_table(headers, rows)

        except Exception as e:
            logging.error(f"Error formatting query result: {e}")
            return "Error formatting query result"

    def get_products_str(self, with_link=True):
        r = self.get_products()
        if not r:
            return "No products found"

        try:
            # Prepare headers and rows for the table
            headers = ["App", "Version", "Version Timestamp"]
            rows = [[row['name'], row['version'], row['timestamp']] for row in r]

            # Format the table
            table = self.format_table(headers, rows)

            if with_link:
                table += self.get_link(use="products", text="Product page")

        except Exception as e:
            logging.error(f"Error decoding JSON: {e}")
            table = "Error getting product list"

        logging.info(table)
        return table

    def get_logical_app_version(self, logical_app):
        if not self.product_list:
            self.refresh_data()
            if not self.product_list:
                return "I have issues getting data from Scribe, Sorry"
        for product in self.product_list:
            if product["name"] == logical_app:
                latest = product["version"]
                logging.info(f"Latest version for {logical_app} is {latest}")
                return latest
        return None

    def get_product_vulnerabilities(self, logical_app, logical_app_version=None):
        self.refresh_data()
        if not logical_app_version:
            v = self.get_logical_app_version(logical_app)
            if v:
                logical_app_version = v
            else:
                logging.error(f"Failed to get version for {logical_app}")
                return None, None

        if not self.validate_dataset(self.scribe_vulnerabilities_dataset):
            return None
        url = f"{self.base_url}/dataset/data"
        body = {
            "superset_token": self.superset_token,
            "validate": False,
            "query": {
                "datasource": {
                    "id": self.dataset_ids[self.scribe_vulnerabilities_dataset]["id"],
                    "type": "table"
                },
                "force": False,
                "queries": [
                    {
                        "columns": [
                            "vulnerability_id",
                            "severity",
                            "epssProbability",
                            "targetName",
                            "component_name",
                            "logical_app",
                        ],
                        "filters": [
                            {
                                "col": "logical_app",
                                "op": "==",
                                "val": logical_app
                            },
                            {
                                "col": "logical_app_version",
                                "op": "==",
                                "val": logical_app_version
                            },
                            {
                                "col": "vulnerability_id",
                                "op": "like",
                                "val": "CVE-%"
                            }
                        ],
                        "metrics": [],
                        "orderby": [["severity", False], ["epssProbability", False]],
                        "row_limit": 10
                    }
                ],
                "result_format": "json",
                "result_type": "results"
            },
            "validate": False
        }

        r = self.send_post_request(url=url, token=self.jwt_token, body=body)
        if r:
            try:
                return r.json(), logical_app_version
            except Exception as e:
                logging.error(f"Error decoding JSON: {e}")
                return {}, logical_app_version
        return {}, logical_app_version

    def get_product_vulnerabilities_str(self, logical_app, logical_app_version=None, with_link=True):
        r, logical_app_version = self.get_product_vulnerabilities(logical_app, logical_app_version)
        if not r:
            return "Failed getting vulnerability data"

        try:
            vuln_list = r["result"][0]["data"]
            if not vuln_list:
                return "No vulnerabilities found"

            # Prepare headers and rows for the table
            headers = ["ID", "Severity", "EPSS", "Vulnerable Component", "Artifact"]
            rows = [
                [
                    f"[{vuln['vulnerability_id']}](https://nvd.nist.gov/vuln/detail/{vuln['vulnerability_id']})",
                    vuln['severity'],
                    vuln['epssProbability'],
                    vuln['component_name'],
                    vuln['targetName']
                ]
                for vuln in vuln_list
            ]

            # Format the table
            table = self.format_table(headers, rows)

            if with_link:
                table += self.get_link(
                    use="product_vulnerabilities",
                    text="Product vulnerabilities",
                    params={"product": logical_app, "product_version": logical_app_version, "show_file_components": "false"}
                )

        except Exception as e:
            logging.error(f"Error decoding JSON: {e}")
            table = "Error getting vulnerability list"

        logging.info(table)
        return table

    def get_product_vulnerability_distribution(self, logical_app, logical_app_version=None):
        def severity_label(severity):
            print(f"Assigning severity label for severity: {severity}")
            if severity <= 3.9:
                return "Low"
            elif severity <= 6.9:
                return "Medium"
            elif severity <= 8.9:
                return "High"
            else:
                return "Critical"
            
        self.refresh_data()
        if not logical_app_version:
            v = self.get_logical_app_version(logical_app)
            if v:
                logical_app_version = v
            else:
                return None

        if not self.validate_dataset(self.scribe_vulnerabilities_dataset):
            return None
        url = f"{self.base_url}/dataset/data"

        body = {
            "superset_token": self.superset_token,
            "validate": False,
            "query": {
                "datasource": {
                    "id": self.dataset_ids[self.scribe_vulnerabilities_dataset]["id"],
                    "type": "table"
                },
                "force": False,
                "queries": [
                    {
                        "columns": [
                            "logical_app",
                            "logical_app_version",
                            "vulnerability_id",
                            "severity"
                        ],
                        "filters": [
                            {
                                "col": "logical_app",
                                "op": "==",
                                "val": logical_app
                            },
                            {
                                "col": "logical_app_version",
                                "op": "==",
                                "val": logical_app_version
                            },
                            {
                                "col": "vulnerability_id",
                                "op": "like",
                                "val": "CVE-%"
                            }
                        ],
                        "metrics": [
                            {
                                "label": "vulnerabilities",
                                "expressionType": "SQL",
                                "sqlExpression": "COUNT(DISTINCT vulnerability_id)"
                            },
                        ],
                        "post_processing": [],
                        "groupby": ["severity"],
                        "orderby": [["severity", False]],
                        "row_limit": 10
                    }
                ],
                "result_format": "json",
                "result_type": "results"
            }
        }

        r = self.send_post_request(url=url, token=self.jwt_token, body=body)
        if r:
            try:
                r = r.json()
                data = r["result"][0]["data"]
                for row in data:
                    row["severity_label"] = severity_label(row["severity"])
                return data
            except Exception as e:
                logging.error(f"Error decoding JSON: {e}")
                return []
        return []


    def get_datasets(self):
        url = f"{self.base_url}/dataset"
        body = {
            "superset_token": self.superset_token
        }

        r = self.send_post_request(url=url, token=self.jwt_token, body=body)
        if r:
            r = r.json()
            os.makedirs("tmp", exist_ok=True)  # Ensure the tmp directory exists

            with open("tmp/datasets.json", "w") as f:
                json.dump(r, f)
            
        return r
    
    def query_superset(self, dataset, query, time_tokens=[]):
        dataset_name = dataset
        id = None
        if self.dataset_ids and dataset in self.dataset_ids:
            id = self.dataset_ids[dataset]["id"]
        else:
            logging.error(f"Dataset {dataset} not found in dataset IDs")
            return None
        
        url = f"{self.base_url}/dataset/data"
        body = {
            "superset_token": self.superset_token,
            "validate": False, 
            "query": {
                "datasource": {
                    "id": id,
                    "type": "table"
                },
                "force": "false",
                "queries": [
                    query
                ],
                "result_format": "json",
                "result_type": "results"
            },
            "validate": False
        }

        r = self.send_post_request(url=url, token=self.jwt_token, body=body)
        if r:
            try:
                r = r.json()
                for t in time_tokens:
                    r = convert_timestamps(r, col_substr=t)
                return r
            except Exception as e:
                logging.error(f"Error decoding JSON: {e}")
                return {}
            
    def query_result_to_str(self, header, r):
        if not header:
            header = ""
        try:
            data = r["result"][0]["data"]
            if not data:
                return "No data found"
            columns = r["result"][0]["colnames"]
            if not columns:
                return "No columns information found"
            
            md_table = ["| " + " | ".join(columns) + " |", "| ---" * len(columns) + " |"]
            for row in data:
                md_table.append("| " + " | ".join([str(row[col]) for col in columns]) + " |")
            o = header + '\n' + "\n".join(md_table)
        except Exception as e:
            logging.error(f"Error decoding JSON: {e}")
            o = "Error getting data"

        return o
            
    def query_str(self, dataset, query, header = "Query results:"):
        r = self.query_superset(dataset, query, time_tokens = ["timestamp", "time_evaluated", "Time", "created", "published_on", "LastModified",
                                                               "date_changed", "advisory_modified"])
        if not r:
            return "Failed getting data"
        o = self.query_result_to_str(header,r)
        return o
    
    def get_policy_results(self, logical_app="Scribot", logical_app_version=None, initiative="%"):
        self.refresh_data()
        if not logical_app_version:
            v = self.get_logical_app_version(logical_app)
            if v:
                logical_app_version = v
            else:
                logging.error(f"Failed to get version for {logical_app}")
                return None, None
        if not self.validate_dataset(self.scribe_policy_dataset):
            return None, None

        url = f"{self.base_url}/dataset/data"
        body = {
            "superset_token": self.superset_token,
            "validate": False,
            "query": {
                "datasource": {
                    "id": self.dataset_ids[self.scribe_policy_dataset]["id"],
                    "type": "table"
                },
                "force": "false",
                "queries": [
                    {
                        "columns": [
                            "logical_app",
                            "logical_app_version",
                            "asset_type",
                            "asset_name",
                            "status",
                            "time_evaluated",
                            "initiative_name",
                            "control_name",
                            "rule_name",
                            "message"
                        ],
                        "filters": [
                            # {
                            #     "col": "initiative_name",
                            #     "op": "like",
                            #     "val": initiative
                            # },
                            {
                                "col": "logical_app",
                                "op": "==",
                                "val": logical_app
                            },
                            {
                                "col": "logical_app_version",
                                "op": "==",
                                "val": logical_app_version
                            },
                            {
                                "col": "status",
                                "op": "not in",
                                "val": ['Not Applicable', 'Not Applicable']
                            }
                        ],
                        "metrics": [],
                        "orderby": [],#["time_evaluated", False], ["status", True]],
                        "row_limit": 10
                    }
                ],
                "result_format": "json",
                "result_type": "results"
            }
        }

        r = self.send_post_request(url=url, token=self.jwt_token, body=body)
        if r:
            try:
                r = r.json()
                r = convert_timestamps(r, "time_evaluated")
                return r, logical_app_version
            except Exception as e:
                logging.error(f"Error decoding JSON: {e}")
                return {}, logical_app_version
        return {}, logical_app_version

    def get_policy_results_str(self, logical_app="scribe-platform", logical_app_version=None, initiative="%", with_link = True):
        r, logical_app_version = self.get_policy_results(logical_app, logical_app_version, initiative)
        if not r:
            return "Failed getting policy results"

        o = f"Policy results for {logical_app}:\n"
        try:
            o = self.query_result_to_str(o, r)

        except Exception as e:
            logging.error(f"Error decoding JSON: {e}")
            o = "Error getting policy results"
        return o
    
     
    def get_product_lineage(self, logical_app, logical_app_version=None):
        self.refresh_data()
        if not logical_app_version:
            v = self.get_logical_app_version(logical_app)
            if v:
                logical_app_version = v
            else:
                logging.error(f"Failed to get version for {logical_app}")
                return None, None
        if not self.validate_dataset(self.scribe_policy_dataset):
            return None, None
        # left out columns:
        # "product_id",
        # "version_id",
        # "child_id",
        # "properties",
        query = {
            "columns": [
                "timestamp",
                "logical_app",
                "logical_app_version",
                "platform_name",
                "platform_type",
                "asset_type",
                "asset_name",
                "external_id",
                "uri",
                "owner",
                "path",
                "parent_id",
                "parent_type",
                "parent_name",
                "parent_external_id",
                # "properties"
            ],
            "filters": [
                {
                    "col": "logical_app",
                    "op": "==",
                    "val": logical_app
                },
                {
                    "col": "logical_app_version",
                    "op": "==",
                    "val": logical_app_version
                }
            ],
            "metrics": [],
            "orderby": [["timestamp", False]],
            "annotation_layers": [],
            "row_limit": 100,
            "post_processing": []
        }
        r = self.query_superset(self.scribe_lineage_dataset, query, time_tokens = ["timestamp"])
        # TODO: check in superset prod for more columns
        if r:
            try:
                data = r["result"][0]["data"]
            except Exception as e:
                logging.error(f"Error decoding JSON: {e}")
                return [], logical_app_version
            if not data:
                return [], logical_app_version
            for row in data:
                if row.get("properties"):
                    try:
                        row["properties"] = ast.literal_eval(row["properties"])
                    except Exception as e:
                        logging.error(f"Error decoding properties: {e}")
                        row["properties"] = {}
            return data, logical_app_version
        else:
            return [], logical_app_version

    types_to_position = {
            "organization": 1,
            "repo": 2,
            "branch": 3,
            "workflow": 4,
            "workflow_run": 5,
            "image": 6,
            "pod": 7,
            "namespace": 8

    }

    
    include_fields = [
        "platform_name", "platform_type",
        "asset_name", "asset_type","parent_name", "parent_type"
   ]

    def filter_lineage_data(self, data, include_types=types_to_position, include_fields=include_fields):
        """
        Filter the data to only include records of the specified types
        and remove specified fields.

        Parameters:
            data (list): List of JSON records.
            include_types (set): Asset types to keep.
            include_fields (set): Fields to remove from records.

        Returns:
            list: Filtered data.
        """
        filtered_data = []
        
        for record in data:
            if record.get("asset_type") in include_types:
                # Create a copy to avoid modifying original data
                filtered_record = {key: value for key, value in record.items() if key in include_fields}
                filtered_data.append(filtered_record)
        
        return filtered_data
    
    def list_attestations(self, criteria={}):
        self.refresh_data()
        url = f"{self.base_url}/evidence/list"
        body = criteria
        r = self.send_post_request(url=url, token=self.jwt_token, body=body)
        if r:
            try:
                return r.json()
            except Exception as e:
                logging.error(f"Error decoding JSON: {e}")
        return {}

    def get_attestation(self, attestation_id):
        url = f"{self.base_url}/evidence/{attestation_id}"
        r = self.send_get_request(url=url, token=self.jwt_token)
        if not r:
            return {}
        try:
            o = r.json()
            url = o.get("presigned_url")
            if url:
                response = urllib.request.urlopen(url)
                content = response.read()
                o = json.loads(content)
            
                if 'payload' in o:
                    o = self.att2statement(o)
            else:
                logging.error(f"Presigned URL not found in response")
                return {}
        except Exception as e:
            logging.error(f"Error decoding JSON: {e}")
            return {}
        return o

    def get_latest_attestation(self, criteria={}):
        atts = self.list_attestations(criteria=criteria)
        if not atts:
            return {}
        evidence_list = atts.get("evidences", [])
        if not evidence_list:
            return {}

        evidence_list.sort(key=lambda x: x['context']["timestamp"], reverse=True)
        if not evidence_list:
            return {}
        evidence = evidence_list[0]

        att = self.get_attestation(evidence["id"])
        return att
    
    def query_vulnerabilities(self, querystr, title="") -> str:
        logging.debug(f"Query vulnerabilities with: {querystr}")
        query = self.querystr_to_json(querystr)
        if query:
            r = self.query_superset(self.scribe_vulnerabilities_dataset, query, time_tokens=["timestamp"])
            if r:
                return f"{title}\n{self.format_query_result_table(r)}"
        return "Error running query"

    def query_products(self, querystr, title="") -> str:
        logging.debug(f"Query products with: {querystr}")
        query = self.querystr_to_json(querystr)
        if query:
            r = self.query_superset(self.scribe_products_dataset, query, time_tokens=["timestamp"])
            if r:
                return f"{title}\n{self.format_query_result_table(r)}"
        return "Error running query"

    def get_lineage_graph(self, lineage_data, output_file):
        """
        Generate a lineage graph from the lineage data and save it to the specified output file.
        
        Parameters:
            lineage_data (list): List of lineage records.
            output_file (str): Path to save the generated graph.
        
        Returns:
            str: Full path of the saved graph file.
        """
        if not lineage_data:
            logging.error("No lineage data provided for graph generation.")
            return "Error: No lineage data provided."
        
        full_path_filename = create_lineage_graph(lineage_data, output_file=output_file)
        return full_path_filename
    
    def query_lineage(self, querystr, title="", graph_filename=None) -> str:
        logging.debug(f"Query lineage with: {querystr}")
        query = self.querystr_to_json(querystr)
        if query:
            r = self.query_superset(self.scribe_lineage_dataset, query, time_tokens=["timestamp"])
            if r:
                if graph_filename:
                    required_fields = ["asset_name", "asset_type", "parent_name", "parent_type", "external_id", "parent_external_id", "uri"]
                    columns = query.get("columns", [])
                    if not all(field in columns for field in required_fields):
                        logging.error(f"Query must include the following fields: {', '.join(required_fields)}")
                        return "Error: Query must include required fields for lineage graph generation."
                    
                    lineage_data = r["result"][0]["data"]
                    full_path_filename = create_lineage_graph(lineage_data, output_file=graph_filename)
                    logging.info(f"Lineage graph saved to {full_path_filename}")
                else:
                    logging.warning("No graph filename provided, lineage graph will not be generated.")
                return f"{title}\n{self.format_query_result_table(r)}"
        return "Error running query"
    
    def query_policy_results(self, querystr, title="") -> str:
        logging.debug(f"Query policy results with: {querystr}")
        query = self.querystr_to_json(querystr)
        if query:
            r = self.query_superset(self.scribe_policy_dataset, query, time_tokens=["time_evaluated"])
            if r:
                return f"{title}\n{self.format_query_result_table(r)}"
        return "Error running query"
    
    def query_risk(self, querystr, title="") -> str:
        logging.debug(f"Query risk table with: {querystr}")
        query = self.querystr_to_json(querystr)
        if query:
            r = self.query_superset(self.scribe_risk_dataset, query, time_tokens=["timestamp"])
            if r:
                return f"{title}\n{self.format_query_result_table(r)}"
        return "Error running query"
    
    def query_findings(self, querystr, title="") -> str:
        logging.debug(f"Query finding table with: {querystr}")
        query = self.querystr_to_json(querystr)
        if query:
            r = self.query_superset(self.scribe_findings_dataset, query, time_tokens=["timestamp"])
            if r:
                return f"{title}\n{self.format_query_result_table(r)}"
        return "Error running query"
    
    def superset_authenticate(self):
        if not self.username or not self.password:
            logging.error("Username and password must be provided for Superset login.")
            return False

        login_url = f"{self.superset_base_url}/api/v1/security/login"
        body = {
            "username": self.username,
            "password": self.password,
            "provider": "db",
            "refresh": True
        }

        response = self.send_post_request(login_url, token = None, body=body)
        if not response:
            logging.error("Login failed. No response received.")
            return False

        try:
            result = response.json()
        except json.JSONDecodeError:
            logging.error("Failed to parse login response.")
            return False

        access_token = result.get("access_token")
        if access_token:
            self.jwt_token = access_token
            self.superset_token = access_token
            logging.info("Login successful. Tokens obtained.")
            return True
        else:
            logging.error("Login failed. No access token found.")
            return False
    
    def get_dataset_links(self):
                
        self.superset_authenticate()

        if not self.jwt_token:
            logging.error("No JWT token found. Please authenticate first.")
            return []

        dataset_links = []

        #Use slug to directly fetch the API dashboard's datasets
        dashboard_slug = "api-dashboard"
        datasets_url = f"{self.superset_base_url}/api/v1/dashboard/{dashboard_slug}/datasets"
        logging.info(f"Fetching datasets from {datasets_url}")

        datasets_response = self.send_get_request(datasets_url, token=self.jwt_token)
        if not datasets_response:
            logging.error("Failed to fetch datasets for API dashboard.")
            return []

        try:
            datasets_data = datasets_response.json()
            datasets = datasets_data.get("result", [])

            for dataset in datasets:
                dataset_id = dataset.get("id")
                if dataset_id:
                    link = f"{self.base_url}/explore/?datasource_type=table&datasource_id={dataset_id}"
                    dataset_links.append(link)

        except json.JSONDecodeError as e:
            logging.error(f"Error parsing dataset response: {e}")

        logging.info(f"Total dataset links collected: {len(dataset_links)}")
        return dataset_links


    def get_dataset_data(self, dataset_links, output_file="dataset_metadata.csv"):
        records = []

        for link in dataset_links:
            #Properly parse dataset ID from the URL
            parsed_url = urlparse(link)
            query_params = parse_qs(parsed_url.query)
            dataset_id_list = query_params.get("datasource_id")
            if not dataset_id_list:
                logging.error(f"Could not find dataset_id in link: {link}")
                continue
            dataset_id = dataset_id_list[0]

            dataset_url = f"{self.superset_base_url}/api/v1/dataset/{dataset_id}"
            logging.info(f"Fetching metadata from {dataset_url}")

            response = self.send_get_request(dataset_url, token=self.jwt_token)
            if not response:
                logging.error(f"Failed to fetch metadata for dataset ID {dataset_id}")
                continue
            try:
                metadata_json = response.json()
                dataset_metadata = metadata_json.get("result", {})
            except json.JSONDecodeError as e:
                logging.error(f"Failed to parse metadata JSON for dataset ID {dataset_id}: {e}")
                continue

            verbose_name = dataset_metadata.get("verbose_name")
            if not verbose_name or verbose_name == dataset_metadata.get("table_name"):
                verbose_name = dataset_metadata.get("table_name", "").replace("_", " ").title()

            #Get dataset metadata
            dataset_info = {
                "dataset_id": dataset_metadata.get("id"),
                "catalog": dataset_metadata.get("catalog"),
                "changed_on": dataset_metadata.get("changed_on"),
                "table_name": dataset_metadata.get("table_name"),
                "verbose_name": verbose_name,
                "schema": dataset_metadata.get("schema"),
                "datasource_type": dataset_metadata.get("datasource_type"),
                "description": dataset_metadata.get("description"),
            }

            #Get column metadata
            columns = dataset_metadata.get("columns", [])
            if columns:
                for col in columns:
                    record = dataset_info.copy()
                    record.update({
                        "column_id": col.get("id"),
                        "column_name": col.get("column_name"),
                        "column_verbose_name": col.get("verbose_name") or col.get("column_name", "").replace("_", " ").title(),
                        "column_type": col.get("type"),
                        "groupby": col.get("groupby"),
                        "filterable": col.get("filterable"),
                        "is_dttm": col.get("is_dttm"),
                        "column_description": col.get("description"),
                    })
                    records.append(record)
            else:
                records.append(dataset_info)

        #Write to CSV
        if records:
            fieldnames = [
                "dataset_id", "table_name", "verbose_name", "schema", "catalog",
                "datasource_type", "changed_on", "description",
                "column_id", "column_name", "column_verbose_name", "column_type",
                "groupby", "filterable", "is_dttm", "column_description"
            ]
            fieldnames = [field for field in fieldnames if any(field in rec for rec in records)]

            with open(output_file, mode='w', newline='', encoding='utf-8') as csvfile:
                writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
                writer.writeheader()
                writer.writerows(records)

            logging.info(f"Dataset metadata exported to {output_file}")
        else:
            logging.warning("No dataset metadata exported.")

    def get_dataset_tables_dict(self, tables=[]) -> dict:
        if not self.jwt_token:
            self.superset_authenticate()

        dataset_links = self.get_dataset_links()
        if not dataset_links:
            logging.error("No dataset links found.")
            return {}

        self.get_dataset_data(dataset_links)

        # Load data from the CSV file (generated by get_dataset_data)
        records = []
        try:
            with open("dataset_metadata.csv", mode="r", encoding="utf-8") as f:
                reader = csv.DictReader(f)
                for row in reader:
                    records.append(row)
        except Exception as e:
            logging.error(f"Failed to load dataset_metadata.csv: {e}")
            return {}

        # Group column data by table name
        table_dict = {}
        for row in records:
            table_name = row.get("table_name")
            if tables and table_name not in tables:
                continue

            if table_name not in table_dict:
                table_dict[table_name] = []

            table_dict[table_name].append({
                "column_name": row.get("column_name", ""),
                "type": row.get("column_type", ""),
                "label": row.get("column_verbose_name") or row.get("column_name", "").replace("_", " ").title(),
                "description": row.get("column_description", "")
            })

        # Add markdown for each table
        for table_name, cols in table_dict.items():
            md_parts = []
            md_parts.append(f"## `{table_name}` Columns\n\n")
            md_parts.append("| Column Name | Label | Type | Description | Like-compliant |\n")
            md_parts.append("|-------------|-------|------|-------------|----------------|\n")
            for col in cols:
                column_name = col["column_name"]
                label = col["label"]
                col_type = col["type"].upper()
                description = col["description"]
                expr = row.get("expression", "") or ""
                if col_type == "STRING" and (" in (" in expr.lower() or " in(" in expr.lower()):
                    like_compliant = False
                else:
                    like_compliant = True
                md_parts.append(
                    f"| `{column_name}` | {label} | {col_type} | {description} | {str(like_compliant)} |\n"
                )
            md_parts.append("\n")
            md_str = "".join(md_parts)
            # Add markdown string to the dict for this table
            # If table_dict[table_name] is a list, convert to dict with 'columns' and 'md'
            table_dict[table_name] = {
                "columns": cols,
                "md": md_str
            }
        return table_dict


    def get_dataset_tables_md(self, tables=[]) -> str:
        if not self.jwt_token:
            self.superset_authenticate()

        dataset_links = self.get_dataset_links()
        if not dataset_links:
            logging.error("No dataset links found.")
            return ""

        self.get_dataset_data(dataset_links)

        # Load data from the CSV file
        records = []
        try:
            with open("dataset_metadata.csv", mode="r", encoding="utf-8") as f:
                reader = csv.DictReader(f)
                for row in reader:
                    records.append(row)
        except Exception as e:
            logging.error(f"Failed to load dataset_metadata.csv: {e}")
            return ""

        # Group columns by table_name
        table_columns = {}
        for row in records:
            table_name = row.get("table_name")
            if tables and table_name not in tables:
                continue

            if table_name not in table_columns:
                table_columns[table_name] = []

            column_name = row.get("column_name", "")
            col_type = row.get("column_type", "")
            col_type = row.get("column_type", "").upper()
            expr = row.get("expression", "") or ""

            if col_type == "STRING" and (" in (" in expr.lower() or " in(" in expr.lower()):
                like_compliant = False
            else:
                like_compliant = True


            table_columns[table_name].append({
                "column_name": column_name,
                "label": row.get("column_verbose_name") or column_name.replace("_", " ").title(),
                "type": col_type,
                "description": row.get("column_description", ""),
                "like_compliant": str(like_compliant)
            })

        # Build markdown
        md_parts = []
        for table_name, cols in table_columns.items():
            md_parts.append(f"## `{table_name}` Columns\n\n")
            md_parts.append("| Column Name | Label | Type | Description | Like-compliant |\n")
            md_parts.append("|-------------|-------|------|-------------|----------------|\n")

            for col in cols:
                md_parts.append(
                    f"| `{col['column_name']}` | {col['label']} | {col['type']} | "
                    f"{col['description']} | {col['like_compliant']} |\n"
                )
            md_parts.append("\n")

        markdown_output = "".join(md_parts)

        # Write to file
        with open("dataset_tables.md", "w", encoding="utf-8") as f:
            f.write(markdown_output)

        logging.info("Saved table metadata Markdown to dataset_tables.md")
        return markdown_output
    
    def get_dataset_prompt(self, dataset, prompt_template=None):
        """
        Returns a prompt string for the given dataset, using a template.
        - dataset: table name to fetch from docs/dataset_tables.json
        - prompt_template: template string with '{table}' placeholder, or None to load from docs/<dataset>-template.md
        """
        import os
        import json
        # Load table data from docs/dataset_tables.json
        table_json_path = os.path.join("docs", "dataset_tables.json")
        if not os.path.exists(table_json_path):
            raise FileNotFoundError(f"dataset_tables.json not found at {table_json_path}")
        with open(table_json_path, "r", encoding="utf-8") as f:
            tables = json.load(f)
        if dataset not in tables:
            raise ValueError(f"Dataset '{dataset}' not found in dataset_tables.json")
        table_md = tables[dataset].get("md")
        if not table_md:
            raise ValueError(f"No markdown found for dataset '{dataset}' in dataset_tables.json")
        # Load template
        if prompt_template is None:
            template_filename = f"{dataset.replace(' ', '_')}-template.md"
            template_path = os.path.join("docs", template_filename)
            if os.path.exists(template_path):
                with open(template_path, "r", encoding="utf-8") as f:
                    prompt_template = f.read()
            else:
                prompt_template = "{table}"
        # Build prompt
        prompt = prompt_template.replace("{table}", table_md)
        return prompt