import re
import zipfile
from collections import defaultdict
from pathlib import Path

import pandas as pd
import yaml

from qubx.core.metrics import TradingSessionResult, _pfl_metrics_prepare
from qubx.utils.misc import blue, cyan, green, magenta, red, yellow


class BacktestsResultsManager:
    """
    Manager class for handling backtesting results.

    This class provides functionality to load, list and manage backtesting results stored in zip files.
    Each result contains trading session information and metrics that can be loaded and analyzed.

    Parameters
    ----------
    path : str
        Path to directory containing backtesting result zip files

    Methods
    -------
    reload()
        Reloads all backtesting results from the specified path
    list(regex="", with_metrics=False)
        Lists all backtesting results, optionally filtered by regex and including metrics
    load(name)
        Loads a specific backtesting result by name
    """

    def __init__(self, path: str):
        self.path = path
        self.reload()

    def reload(self) -> "BacktestsResultsManager":
        self.results = {}
        names = defaultdict(lambda: 0)
        for p in Path(self.path).glob("**/*.zip"):
            with zipfile.ZipFile(p, "r") as zip_ref:
                try:
                    info = yaml.safe_load(zip_ref.read("info.yml"))
                    info["path"] = str(p)
                    n = info.get("name", "")
                    _new_name = n if names[n] == 0 else f"{n}.{names[n]}"
                    names[n] += 1
                    info["name"] = _new_name
                    self.results[_new_name] = info
                except Exception:
                    pass

        # - reindex
        _idx = 1
        for n in sorted(self.results.keys()):
            self.results[n]["idx"] = _idx
            _idx += 1

        return self

    def __getitem__(
        self, name: str | int | list[int] | list[str] | slice
    ) -> TradingSessionResult | list[TradingSessionResult]:
        return self.load(name)

    def load(
        self, name: str | int | list[int] | list[str] | slice
    ) -> TradingSessionResult | list[TradingSessionResult]:
        match name:
            case list():
                return [self.load(i) for i in name]
            case slice():
                return [self.load(i) for i in range(name.start, name.stop, name.step if name.step else 1)]

        for info in self.results.values():
            match name:
                case int():
                    if info.get("idx", -1) == name:
                        return TradingSessionResult.from_file(info["path"])
                case str():
                    if info.get("name", "") == name:
                        return TradingSessionResult.from_file(info["path"])
        raise ValueError(f"No result found for {name}")

    def delete(self, name: str | int | list[int] | list[str] | slice):
        def _del_idx(idx):
            for info in self.results.values():
                if info.get("idx", -1) == idx:
                    Path(info["path"]).unlink()
                    return info.get("name", idx)
            return None

        match name:
            case str():
                nms = [_del_idx(i) for i in self._find_indices(name)]
                self.reload()
                print(f" -> Deleted {red(', '.join(nms))} ...")
                return

            case list():
                nms = [_del_idx(i) for i in name]
                self.reload()
                print(f" -> Deleted {red(', '.join(nms))} ...")
                return

            case slice():
                nms = [_del_idx(i) for i in range(name.start, name.stop, name.step if name.step else 1)]
                self.reload()
                print(f" -> Deleted {red(', '.join(nms))} ...")
                return

        for info in self.results.values():
            match name:
                case int():
                    if info.get("idx", -1) == name:
                        Path(info["path"]).unlink()
                        print(f" -> Deleted {red(info.get('name', name))} ...")
                        self.reload()
                        return
                case str():
                    if info.get("name", "") == name:
                        Path(info["path"]).unlink()
                        print(f" -> Deleted {red(info.get('name', name))} ...")
                        self.reload()
                        return
        print(f" -> No results found for {red(name)} !")

    def _find_indices(self, regex: str):
        for n in sorted(self.results.keys()):
            info = self.results[n]
            s_cls = info.get("strategy_class", "").split(".")[-1]

            try:
                if not re.match(regex, n, re.IGNORECASE):
                    if not re.match(regex, s_cls, re.IGNORECASE):
                        continue
            except Exception:
                if regex.lower() != n.lower() and regex.lower() != s_cls.lower():
                    continue

            yield info.get("idx", -1)

    def list(self, regex: str = "", with_metrics=True, params=False, as_table=False):
        """List backtesting results with optional filtering and formatting.

        Args:
            - regex (str, optional): Regular expression pattern to filter results by strategy name or class. Defaults to "".
            - with_metrics (bool, optional): Whether to include performance metrics in output. Defaults to True.
            - params (bool, optional): Whether to display strategy parameters. Defaults to False.
            - as_table (bool, optional): Return results as a pandas DataFrame instead of printing. Defaults to False.

        Returns:
            - Optional[pd.DataFrame]: If as_table=True, returns a DataFrame containing the results sorted by creation time.
            - Otherwise prints formatted results to console.
        """
        _t_rep = []
        for n in sorted(self.results.keys()):
            info = self.results[n]
            s_cls = info.get("strategy_class", "").split(".")[-1]

            if regex:
                if not re.match(regex, n, re.IGNORECASE):
                    if not re.match(regex, s_cls, re.IGNORECASE):
                        continue

            name = info.get("name", "")
            smbs = ", ".join(info.get("symbols", list()))
            start = pd.Timestamp(info.get("start", "")).round("1s")
            stop = pd.Timestamp(info.get("stop", "")).round("1s")
            dscr = info.get("description", "")
            created = pd.Timestamp(info.get("creation_time", "")).round("1s")
            metrics = info.get("performance", {})
            author = info.get("author", "")
            _s = f"{yellow(str(info.get('idx')))} - {red(name)} ::: {magenta(created)} by {cyan(author)}"

            _one_line_dscr = ""
            if dscr:
                dscr = dscr.split("\n")
                for _d in dscr:
                    _s += f"\n\t{magenta('# ' + _d)}"
                    _one_line_dscr += " " + _d

            _s += f"\n\tstrategy: {green(s_cls)}"
            _s += f"\n\tinterval: {blue(start)} - {blue(stop)}"
            _s += f"\n\tcapital: {blue(info.get('capital', ''))} {info.get('base_currency', '')} ({info.get('commissions', '')})"
            _s += f"\n\tinstruments: {blue(smbs)}"
            if params:
                formats = ["{" + f":<{i}" + "}" for i in [50]]
                _p = pd.DataFrame.from_dict(info.get("parameters", {}), orient="index")
                for i in _p.to_string(
                    max_colwidth=30,
                    header=False,
                    formatters=[(lambda x: cyan(fmt.format(str(x)))) for fmt in formats],
                    justify="left",
                ).split("\n"):
                    _s += f"\n\t  |  {yellow(i)}"

            if not as_table:
                print(_s)

            if with_metrics:
                _m_repr = (
                    pd.DataFrame.from_dict(metrics, orient="index")
                    .T[["gain", "cagr", "sharpe", "qr", "max_dd_pct", "mdd_usd", "fees", "execs"]]
                    .astype(float)
                )
                _m_repr = _m_repr.round(3).to_string(index=False)
                _h, _v = _m_repr.split("\n")
                if not as_table:
                    print("\t " + red(_h))
                    print("\t " + cyan(_v))

            if not as_table:
                print()
            else:
                metrics = {
                    m: round(v, 3)
                    for m, v in metrics.items()
                    if m in ["gain", "cagr", "sharpe", "qr", "max_dd_pct", "mdd_usd", "fees", "execs"]
                }
                _t_rep.append(
                    {"Index": info.get("idx", ""), "Strategy": name}
                    | metrics
                    | {
                        "start": start,
                        "stop": stop,
                        "Created": created,
                        "Author": author,
                        "Description": _one_line_dscr,
                    },
                )

        if as_table:
            _df = pd.DataFrame.from_records(_t_rep, index="Index")
            return _df.sort_values(by="Created", ascending=False)
