"""Models for the Folding Studio API."""

import json
import logging
import os
import sys
import time
from io import StringIO
from pathlib import Path
from typing import Any

import gradio as gr
import numpy as np
from folding_studio import single_job_prediction
from folding_studio.client import Client
from folding_studio.commands.experiment import results as get_results
from folding_studio.commands.experiment import status as get_status
from folding_studio.query import Query
from folding_studio.query.boltz import BoltzQuery
from folding_studio.query.chai import ChaiQuery
from folding_studio.query.protenix import ProtenixQuery
from folding_studio_data_models import AF2Parameters, OpenFoldParameters
from folding_studio_data_models.parameters.base import BaseFoldingParameters

from folding_studio_demo.model_fasta_validators import (
    BaseFastaValidator,
    BoltzFastaValidator,
    ChaiFastaValidator,
    ProtenixFastaValidator,
)


class Capturing(list):
    """Capture stdout output."""

    def __enter__(self):
        self._stdout = sys.stdout
        sys.stdout = self._stringio = StringIO()
        return self

    def __exit__(self, *args):
        self.extend(self._stringio.getvalue().splitlines())
        del self._stringio  # free up some memory
        sys.stdout = self._stdout


logger = logging.getLogger(__name__)


class AF3Model:
    model_name = None

    def __init__(self, api_key: str, query: Query, validator: BaseFastaValidator):
        self.api_key = api_key
        self.query = query
        self.validator = validator

    def call(
        self, seq_file: Path | str, output_dir: Path, format_fasta: bool = False
    ) -> None:
        """Predict protein structure from amino acid sequence using AF3 model.

        Args:
            seq_file (Path | str): Path to FASTA file containing amino acid sequence
            output_dir (Path): Path to output directory
            format_description (bool): Whether to format the description of the sequence
        """
        # Validate FASTA format before calling
        is_valid, error_msg = self.check_file_description(seq_file)
        if format_fasta and not is_valid:
            logger.info("Invalid FASTA file format, forcing formatting...")
            self.format_fasta(seq_file)
        elif not is_valid:
            logger.error(error_msg)
            raise gr.Error(error_msg)

        # Create a client using API key
        logger.info("Authenticating client with API key")
        client = Client.from_api_key(api_key=self.api_key)

        # Define query
        query: Query = self.query.from_file(path=seq_file, query_name="gradio")
        query.save_parameters(output_dir)

        logger.info("Payload: %s", query.payload)

        # Send a request
        logger.info(f"Sending {self.model_name} request to Folding Studio API")
        response = client.send_request(
            query, project_code=os.environ["FOLDING_PROJECT_CODE"]
        )

        # Access confidence data
        logger.info("Confidence data: %s", response.confidence_data)

        response.download_results(output_dir=output_dir, force=True, unzip=True)
        logger.info("Results downloaded to %s", output_dir)

    def format_fasta(self, seq_file: Path | str) -> None:
        """Format sequence to FASTA format.

        Args:
            seq_file (Path | str): Path to FASTA file
        """
        formatted_fasta = self.validator.transform_fasta(seq_file)
        with open(seq_file, "w") as f:
            f.write(formatted_fasta)

    def predictions(self, output_dir: Path) -> list[Path]:
        """Get the path to the prediction.

        Args:
            output_dir (Path): Path to output directory

        Returns:
            list[Path]: List of paths to predictions
        """
        raise NotImplementedError()

    def has_prediction(self, output_dir: Path) -> bool:
        """Check if prediction exists in output directory."""
        return len(self.predictions(output_dir)) > 0

    def check_file_description(self, seq_file: Path | str) -> tuple[bool, str | None]:
        """Check if the file description is correct.

        Args:
            seq_file (Path | str): Path to FASTA file

        Returns:
            tuple[bool, str | None]: Tuple containing a boolean indicating if the format is correct and an error message if not
        """

        is_valid, error_msg = self.validator.is_valid_fasta(seq_file)
        if not is_valid:
            return False, error_msg

        return True, None


class ChaiModel(AF3Model):
    model_name = "Chai"

    def __init__(self, api_key: str):
        super().__init__(api_key, ChaiQuery, ChaiFastaValidator())

    def call(
        self, seq_file: Path | str, output_dir: Path, format_fasta: bool = False
    ) -> None:
        """Predict protein structure from amino acid sequence using Chai model.

        Args:
            seq_file (Path | str): Path to FASTA file containing amino acid sequence
            output_dir (Path): Path to output directory
            format_fasta (bool): Whether to format the FASTA file
        """
        super().call(seq_file, output_dir, format_fasta)

    def predictions(self, output_dir: Path) -> dict[Path, dict[str, Any]]:
        """Get the path to the prediction."""
        prediction = next(output_dir.rglob("pred.model_idx_[0-9].cif"), None)
        if prediction is None:
            return {}

        cif_files = {
            int(f.stem.split("model_idx_")[1]): f
            for f in prediction.parent.glob("pred.model_idx_*.cif")
        }

        # Get all npz files and extract their indices
        npz_files = {
            int(f.stem.split("model_idx_")[1]): f
            for f in prediction.parent.glob("scores.model_idx_*.npz")
        }

        # Find common indices and create pairs
        common_indices = sorted(set(cif_files.keys()) & set(npz_files.keys()))

        return {
            idx: {"prediction_path": cif_files[idx], "metrics": np.load(npz_files[idx])}
            for idx in common_indices
        }


class ProtenixModel(AF3Model):
    model_name = "Protenix"

    def __init__(self, api_key: str):
        super().__init__(api_key, ProtenixQuery, ProtenixFastaValidator())

    def call(
        self, seq_file: Path | str, output_dir: Path, format_fasta: bool = False
    ) -> None:
        """Predict protein structure from amino acid sequence using Protenix model.

        Args:
            seq_file (Path | str): Path to FASTA file containing amino acid sequence
            output_dir (Path): Path to output directory
            format_fasta (bool): Whether to format the FASTA file
        """
        super().call(seq_file, output_dir, format_fasta)

    def predictions(self, output_dir: Path) -> list[Path]:
        """Get the path to the prediction."""
        return list(output_dir.rglob("*_model_[0-9].cif"))


class BoltzModel(AF3Model):
    model_name = "Boltz"

    def __init__(self, api_key: str):
        super().__init__(api_key, BoltzQuery, BoltzFastaValidator())

    def call(
        self, seq_file: Path | str, output_dir: Path, format_fasta: bool = False
    ) -> None:
        """Predict protein structure from amino acid sequence using Boltz model.

        Args:
            seq_file (Path | str): Path to FASTA file containing amino acid sequence
            output_dir (Path): Path to output directory
            format_fasta (bool): Whether to format the FASTA file
        """

        super().call(seq_file, output_dir, format_fasta)

    def predictions(self, output_dir: Path) -> list[Path]:
        """Get the path to the prediction."""
        prediction_paths = list(output_dir.rglob("*_model_[0-9].cif"))
        return {
            int(cif_path.stem[-1]): {
                "prediction_path": cif_path,
                "metrics": np.load(list(cif_path.parent.glob("plddt_*.npz"))[0]),
            }
            for cif_path in prediction_paths
        }


class OldModel:
    model_name = None

    def __init__(self, api_key: str):
        self.api_key = api_key

    def call(
        self,
        seq_file: Path | str,
        output_dir: Path,
        parameters: BaseFoldingParameters,
        *args,
        **kwargs,
    ) -> None:
        """Predict protein structure from amino acid sequence using AF2 model.

        Args:
            seq_file (Path | str): Path to FASTA file containing amino acid sequence
            output_dir (Path): Path to output directory
        """
        output = single_job_prediction(
            fasta_file=seq_file,
            parameters=parameters,
        )
        experiment_id = output["message"]["experiment_id"]
        done = False
        while not done:
            with Capturing() as output:
                get_status(experiment_id)
            status = output[0]
            logger.info(f"Experiment {experiment_id} status: {status}")
            if status == "Done":
                done = True
                logger.info("Downloading results")
                get_results(
                    experiment_id,
                    force=True,
                    unzip=True,
                    output=output_dir / "results.zip",
                )
                logger.info("Results downloaded to %s", output_dir)
            else:
                logger.info("Sleeping for 10 seconds")
                time.sleep(10)

    def format_fasta(self, seq_file: Path | str) -> None:
        """Format sequence to FASTA format.

        Args:
            seq_file (Path | str): Path to FASTA file
        """
        return

    def predictions(self, output_dir: Path) -> dict[int, dict[str, Any]]:
        """Get the path to the prediction.

        Args:
            output_dir (Path): Path to output directory

        Returns:
            dict[int, dict[str, Any]]: Dictionary mapping model indices to their prediction paths and metrics
        """
        prediction_paths = list(
            (output_dir / "results").rglob("relaxed_model_[0-9]_ptm_pred_0.pdb")
        )
        metrics_path = output_dir / "results" / "metrics_per_model.json"
        if not metrics_path.exists():
            return {}
        with open(metrics_path, "r") as f:
            metrics = json.load(f)
        output = {
            int(pred_path.stem.split("_")[2]): {
                "prediction_path": pred_path,
                "metrics": metrics[f"model_{int(pred_path.stem.split('_')[2])}_ptm"],
            }
            for pred_path in prediction_paths
        }
        return output

    def has_prediction(self, output_dir: Path) -> bool:
        """Check if prediction exists in output directory."""
        return len(self.predictions(output_dir)) > 0

    def check_file_description(self, seq_file: Path | str) -> tuple[bool, str | None]:
        """Check if the file description is correct.

        Args:
            seq_file (Path | str): Path to FASTA file

        Returns:
            tuple[bool, str | None]: Tuple containing a boolean indicating if the format is correct and an error message if not
        """

        return True, None


class AF2Model(OldModel):
    model_name = "AlphaFold2"

    def call(self, seq_file: Path | str, output_dir: Path, *args, **kwargs) -> None:
        super().call(seq_file, output_dir, AF2Parameters(), *args, **kwargs)


class OpenFoldModel(OldModel):
    model_name = "OpenFold"

    def call(self, seq_file: Path | str, output_dir: Path, *args, **kwargs) -> None:
        super().call(seq_file, output_dir, OpenFoldParameters(), *args, **kwargs)