import os
import pytest
import sqlite3
from pathlib import Path
from Bio import SeqIO
from Bio import AlignIO
from Bio import Phylo
from nbitk.Tools.hmmalign import Hmmalign
from nbitk.Tools.raxml_ng import RaxmlNg
from nbitk.Tools.megatree_loader import MegatreeLoader
from nbitk.Tools.megatree_pruner import MegatreePruner
from nbitk.config import Config
from nbitk.SeqIO.BCDM import BCDMIterator

# Construct relative path to the test data
TEST_DATA_DIR = Path(__file__).parent.parent / "data"
BCDM_TSV_PATH = TEST_DATA_DIR / "BCDM.tsv"
HMM_PATH = TEST_DATA_DIR / "COI-5P.hmm"
SEQIDLIST = TEST_DATA_DIR / "seqidlist.txt"

# Construct the path to the config file
CONFIG_PATH = Path(__file__).parent.parent.parent / 'config' / 'config.yaml'

@pytest.fixture
def temp_fasta(tmp_path):
    """Fixture to create a temporary FASTA file from BCDM TSV."""
    temp_fasta_path = tmp_path / "temp.fasta"
    with open(BCDM_TSV_PATH, "r") as bcdm_file, open(temp_fasta_path, "w") as fasta_file:
        sequences = SeqIO.parse(bcdm_file, "bcdm-tsv")
        SeqIO.write(sequences, fasta_file, "fasta")
    return temp_fasta_path


@pytest.fixture
def config():
    """Fixture to create and load a Config object for each test."""
    cfg = Config()
    cfg.load_config(CONFIG_PATH)
    cfg.set('log_level', 'INFO')
    return cfg


@pytest.fixture
def aligned_fasta(temp_fasta, config, tmp_path):
    """Fixture to create an aligned FASTA file using hmmalign."""
    stockholm_file = tmp_path / "aligned.sto"
    output_file = tmp_path / "aligned.fasta"

    hmmalign_runner = Hmmalign(config)
    hmmalign_runner.set_hmmfile(str(HMM_PATH))
    hmmalign_runner.set_seqfile(str(temp_fasta))
    hmmalign_runner.set_output(str(stockholm_file))
    hmmalign_runner.set_outformat("Stockholm")
    hmmalign_runner.set_trim()

    return_code = hmmalign_runner.run()
    assert return_code == 0, "hmmalign failed to run successfully"

    # Convert Stockholm to FASTA
    with open(stockholm_file, "r") as sto_handle, open(output_file, "w") as fasta_handle:
        alignments = AlignIO.parse(sto_handle, "stockholm")
        AlignIO.write(alignments, fasta_handle, "fasta")

    return output_file


@pytest.fixture
def raxml_tree(aligned_fasta, config, tmp_path):
    """Fixture to create a phylogenetic tree using raxml-ng."""
    output_prefix = tmp_path / "raxml_output"

    raxml_runner = RaxmlNg(config)
    raxml_runner.set_msa(str(aligned_fasta))
    raxml_runner.set_model("GTR+G")
    raxml_runner.set_prefix(str(output_prefix))
    raxml_runner.set_threads(2)  # Use 2 threads for testing
    raxml_runner.set_search()

    return_code = raxml_runner.run()
    assert return_code == 0, "raxml-ng failed to run successfully"

    best_tree_file = f"{output_prefix}.raxml.bestTree"
    assert os.path.exists(best_tree_file), "RAxML-NG best tree file not found"

    return best_tree_file


@pytest.fixture
def megatree_db(raxml_tree, config, tmp_path):
    """Fixture to create a megatree SQLite database using megatree-loader."""
    output_db = tmp_path / "megatree.sqlite"

    megatree_runner = MegatreeLoader(config)
    megatree_runner.set_infile(str(raxml_tree))
    megatree_runner.set_dbfile(str(output_db))

    return_code = megatree_runner.run()
    assert return_code == 0, "megatree-loader failed to run successfully"

    assert os.path.exists(output_db), "megatree-loader output database file not found"

    return output_db


def test_megatree_pruner(megatree_db, config, tmp_path, raxml_tree):
    """Test running megatree-pruner with the megatree database and seqidlist."""
    output_newick = tmp_path / "pruned_tree.newick"

    megatree_pruner = MegatreePruner(config)
    megatree_pruner.set_dbfile(str(megatree_db))
    megatree_pruner.set_infile(str(SEQIDLIST))
    megatree_pruner.set_outfile(str(output_newick))

    return_code = megatree_pruner.run()

    assert return_code == 0, "megatree-pruner failed to run successfully"
    assert os.path.exists(output_newick), "Output Newick file was not created"

    # Parse the original RAxML-NG tree and the pruned tree
    pruned_tree = Phylo.read(output_newick, "newick")
    input_tree = Phylo.read(raxml_tree, "newick")

    assert len(pruned_tree.get_terminals()) == len(input_tree.get_terminals()), "Pruned tree has incorrect number of terminals"


def test_megatree_pruner_validation(config):
    """Test parameter validation in MegatreePruner."""
    megatree_pruner = MegatreePruner(config)

    # Test missing required parameter
    with pytest.raises(ValueError, match="Database file \\(-d\\) is required for megatree-pruner"):
        megatree_pruner.validate_parameters()

    megatree_pruner.set_dbfile("tree_database.sqlite")
    with pytest.raises(ValueError, match="Either file \\(-i\\) or list \\(-l\\) must be provided for megatree-pruner"):
        megatree_pruner.validate_parameters()

    # Test incompatible parameters
    megatree_pruner.set_infile("taxa_list.txt")
    megatree_pruner.set_list(["taxon1", "taxon2"])
    with pytest.raises(ValueError,
                       match="Only one of file \\(-i\\) or list \\(-l\\) can be provided for megatree-pruner"):
        megatree_pruner.validate_parameters()

