import sys
from typing import Iterator
import requests
import json
from loguru import logger
from tqdm import tqdm
import pathlib
import time
import pandas as pd
import random
import string
import re
import threading
import traceback
import socket
from concurrent.futures import ThreadPoolExecutor
from abc import ABC, abstractmethod
from datetime import datetime
from websocket import create_connection, WebSocketConnectionClosedException
from tables import Session, CompanyFinancials, HistoricalFinancialData

logger.remove()
logger.add(
    sink=lambda msg: tqdm.write(msg, end=""),
    level="INFO",
    format="<green>{time:YYYY-MM-DD at HH:mm:ss}</green> | <level>{level: <8}</level> | <cyan>{file.path}</cyan>:<cyan>{line}</cyan> - <level>{message}</level>",
    colorize=True,
    backtrace=True,
    diagnose=True,  # Enhanced traceback formatting
)

# Constants
PATH = pathlib.Path(__file__).parent.absolute()
PATH_TMP = PATH / "tmp"
SESSION_TIMEOUT = 30  # seconds
MAX_RETRY_ATTEMPTS = 5
REQUEST_DELAY = 0.3  # seconds
WS_URL = "wss://data.tradingview.com/socket.io/websocket"
WS_HEADERS = {"Origin": "https://data.tradingview.com"}


class WebSocketManager:
    """
    A context manager for handling WebSocket connections with automatic reconnection
    and proper resource cleanup.
    """

    def __init__(self, url=WS_URL, headers=WS_HEADERS, max_retries=MAX_RETRY_ATTEMPTS):
        self.url = url
        self.headers = headers
        self.ws = None
        self.max_retries = max_retries

    def __enter__(self):
        """Connect to the WebSocket when entering the context."""
        self.connect()
        return self.ws

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Close the WebSocket when exiting the context."""
        self.close()

    def connect(self):
        """
        Establish a WebSocket connection with retry logic.

        @return: The WebSocket connection object
        @raises ConnectionError: If unable to connect after max_retries
        """
        for attempt in range(self.max_retries):
            try:
                self.ws = create_connection(self.url, headers=json.dumps(self.headers))
                return self.ws
            except Exception as e:
                logger.warning(
                    f"WebSocket connection attempt {attempt+1}/{self.max_retries} failed: {str(e)}"
                )
                if attempt < self.max_retries - 1:
                    sleep_time = REQUEST_DELAY * (attempt + 1)
                    time.sleep(sleep_time)
                else:
                    logger.error(
                        f"Failed to connect to WebSocket after {self.max_retries} attempts:\n{traceback.format_exc()}"
                    )
                    raise ConnectionError(
                        f"Failed to connect to WebSocket after {self.max_retries} attempts"
                    ) from e

    def close(self):
        """Safely close the WebSocket connection."""
        if self.ws:
            try:
                self.ws.close()
            except Exception as e:
                logger.warning(f"Error closing WebSocket: {str(e)}")
            finally:
                self.ws = None

    def reconnect(self):
        """Reconnect the WebSocket if it's closed or experiencing issues."""
        self.close()
        return self.connect()


class Bourse(ABC):
    @abstractmethod
    def isin2ticker(self, isin: str):
        pass

    @abstractmethod
    def ticker2isin(self, ticker: str):
        pass


# Create a global thread pool executor
EXECUTOR_TIMEOUT = ThreadPoolExecutor(max_workers=10)


def timeout(func, args=(), kwargs=None, timeout_duration=1, default=None):
    """
    Execute a function with arguments and keyword arguments with a timeout.

    @param func: The function to execute.
    @param args: A tuple of positional arguments to pass to the function.
    @param kwargs: A dictionary of keyword arguments to pass to the function.
    @param timeout_duration: The maximum amount of time to wait for the function to complete.
    @param default: The value to return if the function times out.
    @return: The result of the function or the default value if the function times out.
    """
    if kwargs is None:
        kwargs = {}

    future = EXECUTOR_TIMEOUT.submit(func, *args, **kwargs)
    try:
        return future.result(timeout=timeout_duration)
    except TimeoutError:
        return default


class tradingview(Bourse):
    countries = [
        "america",
        "canada",
        "austria",
        "belgium",
        "lithuania",
        "latvia",
        "luxembourg",
        "switzerland",
        "netherlands",
        "cyprus",
        "czech",
        "norway",
        "poland",
        "germany",
        "portugal",
        "serbia",
        "russia",
        "denmark",
        "estonia",
        "spain",
        "finland",
        "uk",
        "iceland",
        "hungary",
        "greece",
        "france",
        "sweden",
        "romania",
        "slovakia",
        "turkey",
        "italy",
        "uae",
        "morocco",
        "nigeria",
        "bahrain",
        "egypt",
        "qatar",
        "ksa",
        "tunisia",
        "rsa",
        "kenya",
        "israel",
        "kuwait",
        "argentina",
        "mexico",
        "venezuela",
        "peru",
        "brazil",
        "chile",
        "colombia",
        "australia",
        "malaysia",
        "newzealand",
        "bangladesh",
        "philippines",
        "hongkong",
        "pakistan",
        "china",
        "indonesia",
        "singapore",
        "thailand",
        "india",
        "japan",
        "vietnam",
        "taiwan",
        "korea",
        "srilanka",
    ]
    cache = {"tickers": {}, "isin": {}}
    lock = threading.Lock()
    cache_lock = threading.Lock()

    def __init__(self, workers=5) -> None:
        """
        Initializes the tradingview object with specified workers and preloads cache if available.

        @param workers: Number of worker threads to use for parallel processing.
        """
        self.workers = workers
        self.tqdm_places = [False for _ in range(workers)]
        self.i = 0
        self.last_save = 0
        if (PATH_TMP / "tradingview.json").exists():
            with open(PATH_TMP / "tradingview.json") as f:
                self.cache = json.load(f)

    def save(self, force: bool = False):
        """
        Saves the current cache to a temporary file and replaces the old cache file with the new one.
        """
        if not PATH_TMP.exists():
            PATH_TMP.mkdir()
        if force or time.time() - self.last_save > 5:
            with self.lock:
                with open(PATH_TMP / "tradingview_tmp.json", "w") as f:
                    with self.cache_lock:
                        json.dump(self.cache, f)
                (PATH_TMP / "tradingview_tmp.json").replace(
                    PATH_TMP / "tradingview.json"
                )
                self.last_save = time.time()

    def clean_folder(self):
        """
        Cleans the temporary folder of all tradingview data.
        """
        # 1json = 1csv otherwise delete the csv (use the files_json list to avoid spam the fs)
        to_del = 0
        for f in tqdm((PATH_TMP / "tradingview").glob("*.csv")):
            if not (PATH_TMP / "tradingview" / f"{f.stem}.json").exists():
                f.unlink()
                to_del += 1
        print("DELETED", to_del)

    def isin2ticker(self, _isin: str):
        """
        Retrieves the ticker symbol corresponding to the given ISIN.

        @param _isin: The ISIN for which to find the ticker symbol.
        @return: Ticker symbol if found in cache, or None otherwise.
        """
        if _isin in self.cache["isin"]:
            return self.cache["isin"][_isin]["ticker"]
        for ticker, isin in self.cache["tickers"].items():
            if isin == _isin:
                return ticker

    def ticker2isin(self, _ticker: str):
        """
        Retrieves the ISIN corresponding to the given ticker symbol.

        @param _ticker: The ticker symbol for which to find the ISIN.
        @return: ISIN if found in cache, or None otherwise.
        """
        return self.cache["tickers"].get(_ticker)

    def grab_all_isin(self, force=False, clean=False):
        """
        Grabs all ISINs from the tradingview API for each country listed and updates the cache.

        @param force: If True, force update the cache regardless if it exists. Defaults to False.
        @param clean: If True, clean up cache entries for tickers that no longer exist. Defaults to False.
        """
        if not force and not clean and (PATH_TMP / "tradingview.json").exists():
            logger.info("Using existing cache file. Use force=True to refresh.")
            return

        # Create tmp directory if it doesn't exist
        if not PATH_TMP.exists():
            PATH_TMP.mkdir(parents=True, exist_ok=True)

        session = requests.Session()  # Use session for better performance

        if clean:
            logger.info("Cleaning up cache entries for non-existent tickers")
            tmp = set()
            for country in tqdm(self.countries, desc="Fetching tickers by country"):
                try:
                    data = {
                        "columns": [
                            "logoid",
                            "name",
                            "volume",
                            "market_cap_basic",
                            "sector",
                        ],
                        "ignore_unknown_fields": False,
                        "options": {"lang": "en"},
                        "price_conversion": {"to_currency": "usd"},
                        "range": [0, 500_000],
                        "sort": {"sortBy": "market_cap_basic", "sortOrder": "desc"},
                        "symbols": {},
                        "markets": [country],
                        "filter2": {},
                    }

                    for attempt in range(MAX_RETRY_ATTEMPTS):
                        try:
                            response = session.post(
                                f"https://scanner.tradingview.com/{country}/scan?label-product=screener-stock",
                                headers={"Origin": "https://www.tradingview.com"},
                                data=json.dumps(data),
                                timeout=SESSION_TIMEOUT,
                            ).json()

                            for c in response.get("data", []):
                                if c.get("s"):
                                    tmp.add(c["s"])
                            break
                        except (requests.RequestException, json.JSONDecodeError) as e:
                            logger.warning(
                                f"Error fetching data for {country} (attempt {attempt+1}/{MAX_RETRY_ATTEMPTS}): {str(e)}"
                            )
                            if attempt == MAX_RETRY_ATTEMPTS - 1:
                                logger.error(
                                    f"Failed to fetch data for {country} after {MAX_RETRY_ATTEMPTS} attempts:\n{traceback.format_exc()}"
                                )
                            time.sleep(REQUEST_DELAY * (attempt + 1))
                except Exception as e:
                    logger.error(
                        f"Unexpected error processing country {country}:\n{traceback.format_exc()}"
                    )

            # Remove cache tickers not in the list
            deleted = 0
            kept = 0
            logger.info("Cleaning cache of obsolete tickers...")
            for ticker in tqdm(list(self.cache["tickers"]), desc="Cleaning cache"):
                if ticker not in tmp:
                    deleted += 1
                    isin = self.cache["tickers"][ticker]
                    if isin:  # Check if isin exists and is not empty
                        isin_path = PATH_TMP / "tradingview" / f"{isin}.json"
                        csv_path = PATH_TMP / "tradingview" / f"{isin}.csv"

                        if isin_path.exists():
                            isin_path.unlink()
                        if csv_path.exists():
                            csv_path.unlink()

                        if isin in self.cache["isin"]:
                            del self.cache["isin"][isin]

                    del self.cache["tickers"][ticker]
                else:
                    kept += 1

            logger.info(f"Cache cleanup complete. Deleted: {deleted}, Kept: {kept}")
        else:
            # Add new tickers to cache
            logger.info("Adding new tickers to cache...")
            for country in tqdm(self.countries, desc="Updating ticker cache"):
                try:
                    data = {
                        "columns": [
                            "logoid",
                            "name",
                            "volume",
                            "market_cap_basic",
                            "sector",
                        ],
                        "ignore_unknown_fields": False,
                        "options": {"lang": "en"},
                        "price_conversion": {"to_currency": "usd"},
                        "range": [0, 500_000],
                        "sort": {"sortBy": "market_cap_basic", "sortOrder": "desc"},
                        "symbols": {},
                        "markets": [country],
                        "filter2": {},
                    }

                    for attempt in range(MAX_RETRY_ATTEMPTS):
                        try:
                            response = session.post(
                                f"https://scanner.tradingview.com/{country}/scan?label-product=screener-stock",
                                headers={"Origin": "https://www.tradingview.com"},
                                data=json.dumps(data),
                                timeout=SESSION_TIMEOUT,
                            ).json()

                            new_tickers = 0
                            for c in response.get("data", []):
                                if "s" in c and c["s"] not in self.cache["tickers"]:
                                    self.cache["tickers"][c["s"]] = {}
                                    new_tickers += 1

                            logger.debug(
                                f"Added {new_tickers} new tickers for {country}"
                            )
                            break
                        except (requests.RequestException, json.JSONDecodeError) as e:
                            logger.warning(
                                f"Error fetching data for {country} (attempt {attempt+1}/{MAX_RETRY_ATTEMPTS}): {str(e)}"
                            )
                            if attempt == MAX_RETRY_ATTEMPTS - 1:
                                logger.error(
                                    f"Failed to fetch data for {country} after {MAX_RETRY_ATTEMPTS} attempts:\n{traceback.format_exc()}"
                                )
                            time.sleep(REQUEST_DELAY * (attempt + 1))
                except Exception as e:
                    logger.error(
                        f"Unexpected error processing country {country}:\n{traceback.format_exc()}"
                    )

        # Save the updated cache
        self.save(True)

    @staticmethod
    def chunks(lst, n):
        """
        Yield successive n-sized chunks from lst.
        """
        for i in range(0, len(lst), n):
            yield lst[i : i + n]

    def grab_tickers_info(self):
        """
        Grabs additional information for tickers missing details, using a predefined set of markets.
        """
        places = [
            "EURONEXT",
            "NYSE",
            "NASDAQ",
            "AMEX",
            "LSE",
            "XETRA",
            "BME",
            "MIL",
            "any",
        ]
        tickers = set()
        dones = set()
        for t, _ in self.cache["tickers"].items():
            if _:
                dones.add(t)
        for t, _ in self.cache["tickers"].items():
            if not _:
                if "any" in places or any([t.startswith(p + ":") for p in places]):
                    if t not in dones and t not in tickers:
                        tickers.add(t)
        list_tickers = list(tickers)
        random.shuffle(list_tickers)
        iter_tickers = iter(list_tickers)
        with tqdm(total=len(list_tickers)) as pbar:
            threads = []
            for i in range(self.workers):
                threads.append(
                    threading.Thread(
                        target=self.get_tickers_info,
                        args=(iter_tickers, PATH_TMP / "tradingview", pbar),
                    )
                )
                threads[-1].start()
            for thread in threads:
                thread.join()
            self.save(True)

    @staticmethod
    def generateSession():
        """
        Generates a random session ID.

        @return: A string representing a session ID.
        """
        stringLength = 12
        letters = string.ascii_lowercase
        random_string = "".join(random.choice(letters) for i in range(stringLength))
        return "qs_" + random_string

    @staticmethod
    def prependHeader(st):
        """
        Prepends a header to a WebSocket message.

        @param st: The message to prepend with the header.
        @return: The message with the header prepended.
        """
        return "~m~" + str(len(st)) + "~m~" + st

    @staticmethod
    def constructMessage(func, paramList):
        """
        Constructs a message JSON for WebSocket communication.

        @param func: The function name for the message.
        @param paramList: The parameters for the function.
        @return: A JSON string of the message.
        """
        return json.dumps({"m": func, "p": paramList}, separators=(",", ":"))

    def createMessage(self, func, paramList):
        """
        Creates a WebSocket message with a header and a constructed message body.

        @param func: The function name for the message.
        @param paramList: The parameters for the function.
        @return: The full WebSocket message ready for sending.
        """
        return self.prependHeader(self.constructMessage(func, paramList))

    def sendMessage(self, ws, func, args):
        """
        Sends a message over the WebSocket.

        @param ws: The WebSocket connection.
        @param func: The function name of the message.
        @param args: The arguments for the function.
        """
        ws.send(self.createMessage(func, args))

    def create_session(self, ws, time_session):
        """
        Creates a new session if the previous one is older than 30 seconds.

        @param ws: The WebSocket connection.
        @return: The new session ID.
        """
        if time.time() - time_session > 5:
            time_session = time.time()
        session = self.generateSession()
        self.sendMessage(ws, "quote_create_session", [session])
        return session, time_session

    def get_tickers_info(self, tickers: Iterator, path: pathlib.Path, pbar=None):
        with WebSocketManager() as ws:
            time_session = 0

            def ask(TICK, time_session):
                if pbar:
                    pbar.update(1)
                session, time_session = self.create_session(ws, 0)
                self.sendMessage(
                    ws,
                    "quote_add_symbols",
                    [session, TICK, '={"session":"extended","symbol":"%s"}' % TICK],
                )
                self.sendMessage(ws, "quote_remove_symbols", [session, TICK])
                return time_session

            path.mkdir(parents=True, exist_ok=True)
            position = self.tqdm_places.index(False)
            self.tqdm_places[position] = True
            # Initialize the progress bar with a total value of 100 (we'll update dynamically)
            local_pbar = tqdm(total=100, desc=f"Thread {position+1}")
            done = False
            last_update = 0
            pattern = re.compile("~m~\d+~m~~h~\d+$")
            splitter = re.compile("~m~\d+~m~")
            while not done or time.time() - last_update < 5:
                if time.time() - last_update > 0.3:
                    tick = next(tickers, None)
                    if tick:
                        time_session = ask(tick, time_session)
                        last_update = time.time()
                    else:
                        done = True
                result = ws.recv()
                if not result:
                    continue
                if pattern.match(result):
                    ws.send(result)
                for res in splitter.split(result):
                    if not res.startswith("{"):
                        continue
                    try:
                        m = json.loads(res)
                        if "m" in m and m["m"] == "qsd":
                            r = m["p"][1]["v"]
                            if r.get("isin") and r.get("pro_name"):
                                last_update = time.time()
                                with self.cache_lock:
                                    self.cache["tickers"][r["pro_name"]] = r["isin"]
                                # Update our local progress bar
                                local_pbar.update(1)
                                if pbar:  # Update the parent progress bar if provided
                                    pbar.update(1)
                                self.i += 1
                                if self.i % 100 == 0:
                                    self.save()
                            if "earnings_fy_h" in r:
                                last_update = time.time()
                                if (
                                    r["symbol-proname"] in self.cache["tickers"]
                                ) and self.cache["tickers"][r["symbol-proname"]]:
                                    r["isin"] = self.cache["tickers"][
                                        r["symbol-proname"]
                                    ]
                                    self.store_company_data_in_db(r)
                                tick = next(tickers, None)
                                if tick:
                                    time_session = ask(tick, time_session)
                                else:
                                    done = True
                    except Exception as e:
                        logger.error(
                            f"Error processing ticker data:\n{traceback.format_exc()}"
                        )
                        print(e)
            # Close the local progress bar
            local_pbar.close()
            self.tqdm_places[position] = False

    def all_historical_data(self, force=False):
        tickers = {}
        for ticker, isin in self.cache["tickers"].items():
            if not isin:
                continue
            if isin in self.cache["isin"]:
                if force:
                    tickers[isin] = {self.cache["isin"][isin]["ticker"]}
            else:
                if isin in tickers:
                    tickers[isin].add(ticker)
                else:
                    tickers[isin] = {ticker}
        list_tickers = list(tickers.items())
        random.shuffle(list_tickers)
        tickers_iter = iter(list_tickers)
        total = sum((len(t) for t in tickers.values()))
        print("To download:", total)
        threads = []
        bar = tqdm(total=total)
        for i in range(min(self.workers, total)):
            threads.append(
                threading.Thread(
                    target=self.historical_datas,
                    args=(tickers_iter, PATH_TMP / "tradingview", i, bar),
                )
            )
            threads[-1].start()
        for thread in threads:
            thread.join()

    @staticmethod
    def connect():
        for i in range(5):
            try:
                return create_connection(f"wss://data.tradingview.com/socket.io/websocket",headers=json.dumps({'Origin': 'https://data.tradingview.com'}))
            except:
                time.sleep(10)
    
    def historical_datas(self, tickers: Iterator, path: pathlib.Path, i: int, bar=None):
        """
        Retrieves historical data for multiple tickers and saves it to CSV files.

        @param tickers: Iterator of (isin, ticker_set) pairs
        @param path: Path where to save the CSV files
        @param i: Thread identifier for progress reporting
        @param bar: Progress bar to update
        """
        # Create directory if it doesn't exist
        if not path.exists():
            path.mkdir(parents=True, exist_ok=True)

        ws = self.connect()

        for isin, _tickers in tqdm(tickers, position=i + 1, desc=f"Thread {i+1}"):
            if not isin:  # Skip empty ISINs
                logger.warning("Skipping empty ISIN")
                continue

            # Track the best data
            max_len = 0
            best_df = None
            best_ticker = None
            valid = True

            for ticker in _tickers:
                try:
                    # Try up to 2 times to get historical data
                    for attempt in range(2):
                        try:
                            logger.debug(
                                f"Fetching historical data for {ticker} (ISIN: {isin})"
                            )
                            data = self.historical_data(ws, ticker)
                            length = len(data)
                            break
                        except Exception as e:
                            logger.warning(f"Error on attempt {attempt+1}: {str(e)}")
                            ws = self.connect()

                    if length:
                        logger.debug(f"Got {length} data points for {ticker}")
                        if length > max_len:
                            max_len = length
                            best_ticker = ticker
                            best_df = data
                    else:
                        logger.warning(f"No data for {ticker}")
                except Exception as e:
                    logger.error(
                        f"Error processing {ticker}:\n{traceback.format_exc()}"
                    )
                    valid = False
                    break

                if bar:
                    bar.update(1)

            # Save the best data to disk
            if (
                valid
                and best_ticker
                and isinstance(best_df, pd.DataFrame)
                and not best_df.empty
            ):
                try:
                    best_df.to_csv(path / f"{isin}.csv", date_format="%Y-%m-%d")
                    logger.info(
                        f"Saved {max_len} points for {isin} using {best_ticker}"
                    )

                    # Update the cache
                    with self.cache_lock:
                        if isin not in self.cache["isin"]:
                            self.cache["isin"][isin] = {}
                        self.cache["isin"][isin]["ticker"] = best_ticker

                    # Save cache periodically
                    self.save()
                except Exception as e:
                    logger.error(
                        f"Failed to save data for {isin}:\n{traceback.format_exc()}"
                    )

    def historical_data(self, ws, ticker: str) -> pd.DataFrame:
        """
        Retrieve historical price data for a specific ticker symbol with robust connection handling.

        @param ws: WebSocket connection
        @param ticker: Ticker symbol to retrieve data for
        @return: DataFrame with historical price data (date, open, high, low, close, volume)
        @raises: ValueError if unable to retrieve data for the ticker
        """
        chart_session = self.generateSession()
        max_retries = 3

        # Initialize an empty DataFrame as fallback
        empty_df = pd.DataFrame()

        for attempt in range(max_retries):
            try:
                # Check if we need to reconnect
                if attempt > 0:
                    logger.warning(
                        f"Reconnecting WebSocket for {ticker} (attempt {attempt+1}/{max_retries})"
                    )
                    try:
                        ws = WebSocketManager().connect()
                    except Exception as e:
                        logger.error(f"Failed to reconnect WebSocket: {str(e)}")
                        time.sleep(min(5, 2**attempt))  # Exponential backoff
                        continue

                # Create session and setup chart
                chart_session = self.generateSession()
                self.sendMessage(ws, "chart_create_session", [chart_session, ""])
                self.sendMessage(ws, "switch_timezone", [chart_session, "Etc/UTC"])
                self.sendMessage(
                    ws,
                    "resolve_symbol",
                    [
                        chart_session,
                        "sds_sym_1",
                        f'={{"symbol":"{ticker}","adjustment":"splits","currency-id":"EUR"}}',
                    ],
                )
                self.sendMessage(
                    ws,
                    "create_series",
                    [chart_session, "sds_1", "s1", "sds_sym_1", "1D", 50_000],
                )

                bars = []
                end = False
                timeout_start = time.time()
                timeout_limit = SESSION_TIMEOUT
                pattern = re.compile("~m~\d+~m~~h~\d+$")

                # Set a socket timeout to avoid hanging
                ws.settimeout(5)

                while not end and (time.time() - timeout_start < timeout_limit):
                    try:
                        resp = ws.recv()

                        # Handle ping messages
                        if pattern.match(resp):
                            ws.send(resp)  # Respond to ping with the same message
                            continue

                        # Parse response data
                        for r in resp.split("~m~"):
                            if not r.startswith("{"):
                                continue

                            m = json.loads(r)
                            if (
                                "m" in m
                                and m["m"] == "timescale_update"
                                and "p" in m
                                and len(m["p"])
                                and m["p"][0] == chart_session
                            ):
                                # Process historical data bars
                                if "sds_1" in m["p"][1] and "s" in m["p"][1]["sds_1"]:
                                    hist = m["p"][1]["sds_1"]["s"]
                                    for bar in hist:
                                        bar = bar["v"]
                                        bars.append(
                                            {
                                                "date": datetime.utcfromtimestamp(
                                                    bar[0]
                                                ).date(),
                                                "open": bar[1],
                                                "high": bar[2],
                                                "low": bar[3],
                                                "close": bar[4],
                                                "volume": bar[5],
                                            }
                                        )
                                    end = True
                                    break

                    except WebSocketConnectionClosedException:
                        logger.warning(
                            f"WebSocket connection closed during data retrieval for {ticker}"
                        )
                        break  # Break the inner loop to trigger a reconnection attempt

                    except socket.timeout:
                        # Check if we've spent too much time already
                        if time.time() - timeout_start >= timeout_limit * 0.8:
                            logger.warning(
                                f"Socket timeout for {ticker}, trying to reconnect"
                            )
                            break  # Break inner loop to reconnect
                        continue  # Otherwise just continue waiting

                    except json.JSONDecodeError:
                        logger.warning(
                            f"Failed to decode WebSocket response for {ticker}"
                        )
                        continue

                    except Exception as e:
                        logger.error(
                            f"Unexpected error retrieving data for {ticker}: {str(e)}"
                        )
                        break  # Break inner loop to reconnect

                # Check if we got data in this attempt
                if bars:
                    # Process the collected bars into a DataFrame
                    df = pd.DataFrame.from_dict(bars)
                    df["volume"] = df["volume"].astype(int)
                    df.set_index("date", inplace=True)
                    return df

                # If we didn't get data but there are more attempts left, continue to next attempt
                logger.warning(
                    f"No data retrieved for {ticker} on attempt {attempt+1}/{max_retries}"
                )

            except Exception as e:
                logger.error(
                    f"Error in historical_data for {ticker} (attempt {attempt+1}): {str(e)}"
                )
                if attempt < max_retries - 1:
                    time.sleep(1)  # Wait before retrying

        # If we exhausted all retries without success
        logger.warning(
            f"Failed to retrieve historical data for {ticker} after {max_retries} attempts"
        )
        return empty_df

    def store_company_data_in_db(self, company_data):
        """
        Store company financial data in the database instead of writing to a JSON file.

        @param company_data: Dictionary containing company financial data
        @return: True if successful, False otherwise
        """
        if not company_data.get("isin"):
            logger.error("Missing ISIN in company data")
            return False

        try:
            # Create a new session
            database_session = Session()

            try:
                # Check if company exists in database
                company = (
                    database_session.query(CompanyFinancials)
                    .filter_by(isin=company_data["isin"])
                    .first()
                )

                # If not, create a new record
                if not company:
                    company = CompanyFinancials(isin=company_data["isin"])
                    logger.info(
                        f"Creating new database record for {company_data['isin']}"
                    )
                else:
                    logger.info(f"Updating existing record for {company_data['isin']}")

                # Map the data to the company record
                self._map_company_data_to_model(company, company_data)

                # Save historical data if available
                if any(key.endswith("_h") for key in company_data.keys()):
                    self._store_historical_data(database_session, company_data)

                # Add and commit the company record
                database_session.add(company)
                database_session.commit()

                logger.info(
                    f"Successfully stored data for {company_data['isin']} in database"
                )
                return True

            except Exception as e:
                database_session.rollback()
                logger.error(f"Database error: {str(e)}\n{traceback.format_exc()}")
                return False
            finally:
                database_session.close()

        except Exception as e:
            logger.error(
                f"Error storing company data: {str(e)}\n{traceback.format_exc()}"
            )
            return False

    def _map_company_data_to_model(self, company, data):
        """
        Maps data from TradingView JSON format to SQLAlchemy model fields.

        @param company: CompanyFinancials SQLAlchemy model instance
        @param data: Dictionary containing company financial data
        """
        # Basic company information
        if "symbol-proname" in data:
            company.symbol = data.get("symbol-proname")

        # Map common fields directly
        for field in [
            "exchange_listed",
            "country",
            "country_code_fund",
            "region",
            "sector",
            "industry",
            "business_description",
            "website_url",
            "ceo",
            "founded",
            "cik_code",
            "location",
            "market_cap_basic",
            "currency",
            "currency_id",
            "currency_fund",
            "popularity",
            "popularity_rank",
        ]:
            if field in data:
                setattr(company, field, data.get(field))

        # Map nested or specially named fields
        if "local_popularity" in data:
            company.local_popularity = data.get("local_popularity")

        if "local_popularity_rank" in data:
            company.local_popularity_rank = data.get("local_popularity_rank")

        # Market data
        if "market_cap_calc" in data:
            company.market_cap_calc = data.get("market_cap_calc")

        if "total_shares_outstanding" in data:
            company.total_shares_outstanding = data.get("total_shares_outstanding")

        if "float_shares_outstanding" in data:
            company.float_shares_outstanding = data.get("float_shares_outstanding")

        if "float_shares_outstanding_current" in data:
            company.float_shares_outstanding_current = data.get(
                "float_shares_outstanding_current"
            )

        # Pricing multiples
        if "price_earnings" in data:
            company.price_earnings = data.get("price_earnings")

        if "price_book_ratio" in data:
            company.price_book_ratio = data.get("price_book_ratio")

        if "price_sales_ratio" in data:
            company.price_sales_ratio = data.get("price_sales_ratio")

        if "price_revenue_ttm" in data:
            company.price_revenue_ttm = data.get("price_revenue_ttm")

        # Financial metrics
        if "dividends_yield_current" in data:
            company.dividends_yield_current = data.get("dividends_yield_current")

        if "dividend_yield_recent" in data:
            company.dividend_yield_recent = data.get("dividend_yield_recent")

        # Employee information
        if "number_of_employees" in data:
            company.number_of_employees = data.get("number_of_employees")

        if "number_of_shareholders" in data:
            company.number_of_shareholders = data.get("number_of_shareholders")

        # Income statement metrics
        for period in ["ttm", "fq", "fy"]:
            for field in [
                "total_revenue",
                "gross_profit",
                "oper_income",
                "ebit",
                "ebitda",
                "pretax_income",
                "net_income",
                "diluted_net_income",
            ]:
                field_name = f"{field}_{period}"
                if field_name in data:
                    setattr(company, field_name, data.get(field_name))

        # Balance sheet metrics
        for period in ["fq", "fy"]:
            for field in [
                "total_assets",
                "total_liabilities",
                "total_equity",
                "cash_n_equivalents",
                "total_debt",
                "goodwill",
                "intangibles_net",
            ]:
                field_name = f"{field}_{period}"
                if field_name in data:
                    setattr(company, field_name, data.get(field_name))

        # Cash flow metrics
        for period in ["ttm", "fq", "fy"]:
            for field in [
                "cash_f_operating_activities",
                "cash_f_investing_activities",
                "cash_f_financing_activities",
                "free_cash_flow",
                "capital_expenditures",
            ]:
                field_name = f"{field}_{period}"
                if field_name in data:
                    setattr(company, field_name, data.get(field_name))

        # Per share metrics
        for period in ["ttm", "fq", "fy"]:
            for field in [
                "earnings_per_share_diluted",
                "earnings_per_share_basic",
                "book_value_per_share",
                "dps_common_stock_prim_issue",
            ]:
                field_name = f"{field}_{period}"
                if field_name in data:
                    setattr(company, field_name, data.get(field_name))

        # Margin metrics
        for period in ["ttm", "current", "fq", "fy"]:
            for field in [
                "net_margin",
                "operating_margin",
                "gross_margin",
                "ebitda_margin",
            ]:
                field_name = f"{field}_{period}"
                if field_name in data:
                    setattr(company, field_name, data.get(field_name))

        # Return metrics
        for period in ["ttm", "current", "fq", "fy"]:
            for field in [
                "return_on_equity",
                "return_on_assets",
                "return_on_invested_capital",
            ]:
                field_name = f"{field}_{period}"
                if field_name in data:
                    setattr(company, field_name, data.get(field_name))

        # Dividend information
        if "amount_upcoming" in data:
            company.amount_upcoming = data.get("amount_upcoming")

        if "payment_date_upcoming" in data:
            timestamp = data.get("payment_date_upcoming")
            if timestamp:
                company.payment_date_upcoming = datetime.fromtimestamp(timestamp)

        if "ex_dividend_date_upcoming" in data:
            timestamp = data.get("ex_dividend_date_upcoming")
            if timestamp:
                company.ex_dividend_date_upcoming = datetime.fromtimestamp(timestamp)

        if "amount_recent" in data:
            company.amount_recent = data.get("amount_recent")

        if "payment_date_recent" in data:
            timestamp = data.get("payment_date_recent")
            if timestamp:
                company.payment_date_recent = datetime.fromtimestamp(timestamp)

        if "ex_dividend_date_recent" in data:
            timestamp = data.get("ex_dividend_date_recent")
            if timestamp:
                company.ex_dividend_date_recent = datetime.fromtimestamp(timestamp)

        if "dividend_type_h" in data:
            company.dividend_type = data.get("dividend_type_h")

        # Store any serializable JSON fields that have "_h" suffix as JSON
        for key in data:
            if key.endswith("_h") and not hasattr(company, key):
                try:
                    setattr(company, key, data.get(key))
                except Exception as e:
                    logger.warning(f"Could not set {key} on company model: {str(e)}")

    def _store_historical_data(self, session, data):
        """
        Process and store historical financial data from fields ending with "_h".

        @param session: SQLAlchemy session
        @param data: Dictionary containing company data with historical fields
        """
        isin = data.get("isin")
        if not isin:
            logger.warning("Cannot store historical data: missing ISIN")
            return

        # Process earnings history
        if "earnings_fy_h" in data and isinstance(data["earnings_fy_h"], list):
            for entry in data["earnings_fy_h"]:
                try:
                    period_end = entry.get("period_end")
                    period = entry.get("period")

                    if not period_end or not period:
                        continue

                    # Convert timestamp to date
                    period_end_date = datetime.fromtimestamp(period_end).date()

                    # Check if record already exists
                    historical = (
                        session.query(HistoricalFinancialData)
                        .filter_by(
                            company_isin=isin,
                            fiscal_period_label=period,
                            period_type="FY",
                        )
                        .first()
                    )

                    if not historical:
                        historical = HistoricalFinancialData(
                            company_isin=isin,
                            fiscal_period_label=period,
                            period_type="FY",
                            period_end_date=period_end_date,
                        )

                    # Map the data
                    historical.earnings_per_share_diluted = entry.get(
                        "earnings_per_share_diluted"
                    )
                    historical.earnings_per_share_basic = entry.get(
                        "earnings_per_share_basic"
                    )
                    historical.total_revenue = entry.get("revenue")
                    historical.net_income = entry.get("net_income")

                    session.add(historical)

                except Exception as e:
                    logger.error(f"Error processing historical FY data: {str(e)}")

        # Process quarterly earnings history
        if "earnings_fq_h" in data and isinstance(data["earnings_fq_h"], list):
            for entry in data["earnings_fq_h"]:
                try:
                    period_end = entry.get("period_end")
                    period = entry.get("period")
                    period_type = "FQ"

                    # Determine if this is a half-year (FH) record
                    if period and ("H1" in period or "H2" in period):
                        period_type = "FH"

                    if not period_end or not period:
                        continue

                    # Convert timestamp to date
                    period_end_date = datetime.fromtimestamp(period_end).date()

                    # Check if record already exists
                    historical = (
                        session.query(HistoricalFinancialData)
                        .filter_by(
                            company_isin=isin,
                            fiscal_period_label=period,
                            period_type=period_type,
                        )
                        .first()
                    )

                    if not historical:
                        historical = HistoricalFinancialData(
                            company_isin=isin,
                            fiscal_period_label=period,
                            period_type=period_type,
                            period_end_date=period_end_date,
                        )

                    # Map the data
                    historical.earnings_per_share_diluted = entry.get(
                        "earnings_per_share_diluted"
                    )
                    historical.earnings_per_share_basic = entry.get(
                        "earnings_per_share_basic"
                    )
                    historical.total_revenue = entry.get("revenue")
                    historical.net_income = entry.get("net_income")

                    session.add(historical)

                except Exception as e:
                    logger.error(f"Error processing historical FQ/FH data: {str(e)}")


if __name__ == "__main__":
    tradingv = tradingview()
    # tradingv.grab_all_isin(force=True)
    # tradingv.grab_tickers_info()
    tradingv.all_historical_data()