"""
Bioinformatics Protein Analysis Application
Optimized for Hugging Face Spaces & Docker Deployment
"""

import os
import io
import csv
import sys
import subprocess
import pandas as pd
import joblib
import numpy as np
import torch
import esm
import logging
from pathlib import Path
from datetime import datetime
from flask import Flask, request, render_template, send_file, session, redirect, url_for
from Bio.Seq import Seq

# ========== LOGGING SETUP ==========
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s %(message)s",
    handlers=[
        logging.FileHandler("app.log"),
        logging.StreamHandler(sys.stdout)
    ],
)
logger = logging.getLogger(__name__)

# ========== ENVIRONMENT & PATHS ==========
# Get base directory (works in Docker and local)
BASE_DIR = Path(os.getenv("BASE_DIR", "/app"))
DATA_DIR = BASE_DIR / "data"
TEMP_DIR = BASE_DIR / "temp"
MODEL_DIR = BASE_DIR / "models"
TEMPLATE_DIR = BASE_DIR / "templates"
STATIC_DIR = BASE_DIR / "static"

# Create directories if they don't exist
for dir_path in [DATA_DIR, TEMP_DIR, MODEL_DIR]:
    dir_path.mkdir(parents=True, exist_ok=True)

# Model paths - use defaults if not found
MODEL_PATH = MODEL_DIR / "best_model_LR.sav"
ENCODER_PATH = MODEL_DIR / "label_encoder.pkl"

# BLAST configuration
BLAST_PATH = os.getenv("BLAST_PATH", "/usr/bin/blastp")  # Linux default
BLAST_DB = DATA_DIR / "pathway_db"
MAPPING_FILE = DATA_DIR / "pathway_map.csv"

# Thresholds
BITSCORE_THRESHOLD = 80.0
CONFIDENCE_THRESHOLD = 0.7
RESULT_CSV = TEMP_DIR / "combined_results.csv"

# ========== MODEL LOADING ==========
model = None
encoder = None
esm_model = None
batch_converter = None

logger.info("===== Application Startup =====")
logger.info(f"BASE_DIR: {BASE_DIR}")
logger.info(f"DATA_DIR: {DATA_DIR}")
logger.info(f"TEMP_DIR: {TEMP_DIR}")

# Load Logistic Regression Model
try:
    if MODEL_PATH.exists():
        model = joblib.load(MODEL_PATH)
        logger.info("✅ LR model loaded successfully.")
    else:
        logger.warning(f"⚠️ Model not found at {MODEL_PATH}")
except Exception as e:
    logger.error(f"❌ Failed to load LR model: {e}")

# Load Label Encoder
try:
    if ENCODER_PATH.exists():
        encoder = joblib.load(ENCODER_PATH)
        logger.info("✅ Label encoder loaded successfully.")
    else:
        logger.warning(f"⚠️ Encoder not found at {ENCODER_PATH}")
except Exception as e:
    logger.error(f"❌ Failed to load label encoder: {e}")

# Load ESM2 Model (downloads if not cached)
try:
    logger.info("Loading ESM2 model (this may take a moment on first run)...")
    esm_model, alphabet = esm.pretrained.load_model_and_alphabet("esm2_t6_8M_UR50D")
    esm_model.eval()
    esm_model = esm_model.to("cpu")
    batch_converter = alphabet.get_batch_converter()
    logger.info("✅ ESM2 model loaded successfully on CPU.")
except Exception as e:
    logger.error(f"❌ Failed to load ESM2 model: {e}")

# ========== HELPER FUNCTIONS ==========

def get_best_frame_protein(nuc_seq):
    """Translate DNA to best reading frame protein"""
    nuc_seq = "".join(nuc_seq.split()).upper()
    trim_len = len(nuc_seq) - (len(nuc_seq) % 3)
    nuc_seq = nuc_seq[:trim_len]
    
    seq = Seq(nuc_seq)
    frames = [
        seq[i:].translate(to_stop=False) for i in range(3)
    ] + [
        seq.reverse_complement()[i:].translate(to_stop=False) for i in range(3)
    ]
    
    best_protein = ""
    for prot in frames:
        longest_orf = max(prot.split("*"), key=len)
        if len(longest_orf) > len(best_protein):
            best_protein = longest_orf
    
    return str(best_protein)


def detect_and_translate(seq):
    """Detect if DNA or protein and handle accordingly"""
    seq = seq.upper()
    
    # Check if DNA
    if all(base in "ATGCN" for base in seq):
        logger.info("Detected DNA sequence → translating to protein.")
        if len(seq) < 150:
            logger.warning("DNA sequence too short for translation")
            return None
        
        prot_seq = get_best_frame_protein(seq)
        if len(prot_seq) < 50:
            logger.warning("Translated protein too short")
            return None
        return prot_seq
    else:
        logger.info("Detected protein sequence → using directly.")
        if len(seq) < 50:
            logger.warning("Protein sequence too short")
            return None
        return seq


def esm2_320_embed(sequence):
    """Generate ESM2 embeddings for sequence"""
    if esm_model is None or batch_converter is None:
        raise RuntimeError("ESM model not available.")
    
    batch_labels, batch_strs, batch_tokens = batch_converter([("seq1", sequence)])
    
    with torch.no_grad():
        results = esm_model(batch_tokens, repr_layers=[6], return_contacts=False)
    
    token_representations = results["representations"][6]
    seq_repr = token_representations[0, 1:-1].detach().cpu().numpy()
    
    return seq_repr.mean(axis=0)


def run_blast_and_get_dataframe(temp_fasta, blast_output):
    """Execute BLAST and return results as DataFrame"""
    try:
        # Check if BLAST is available
        if not Path(BLAST_PATH).exists():
            logger.warning(f"⚠️ BLAST not found at {BLAST_PATH}, skipping BLAST search")
            return pd.DataFrame(
                columns=["qseqid", "sseqid", "pident", "length", "evalue", "bitscore"]
            )
        
        # Check if database exists
        if not BLAST_DB.exists():
            logger.warning(f"⚠️ BLAST database not found at {BLAST_DB}, skipping BLAST search")
            return pd.DataFrame(
                columns=["qseqid", "sseqid", "pident", "length", "evalue", "bitscore"]
            )
        
        cmd = [
            str(BLAST_PATH),
            "-query", str(temp_fasta),
            "-db", str(BLAST_DB),
            "-out", str(blast_output),
            "-outfmt", "6 qseqid sseqid pident length evalue bitscore"
        ]
        
        logger.info(f"🧬 Running BLAST...")
        proc = subprocess.run(cmd, capture_output=True, text=True, timeout=300)
        
        if proc.returncode != 0:
            logger.error(f"❌ BLAST error: {proc.stderr}")
            return pd.DataFrame(
                columns=["qseqid", "sseqid", "pident", "length", "evalue", "bitscore"]
            )
        
        if not Path(blast_output).exists() or Path(blast_output).stat().st_size == 0:
            logger.warning("⚠️ No BLAST results generated.")
            return pd.DataFrame(
                columns=["qseqid", "sseqid", "pident", "length", "evalue", "bitscore"]
            )
        
        cols = ["qseqid", "sseqid", "pident", "length", "evalue", "bitscore"]
        df = pd.read_csv(blast_output, sep="\t", names=cols)
        df["bitscore"] = df["bitscore"].astype(float)
        logger.info(f"✅ BLAST completed successfully with {len(df)} hits")
        
        return df
    
    except subprocess.TimeoutExpired:
        logger.error("❌ BLAST timed out (>300s)")
        return pd.DataFrame(
            columns=["qseqid", "sseqid", "pident", "length", "evalue", "bitscore"]
        )
    except Exception as e:
        logger.error(f"❌ Unexpected error in BLAST: {e}")
        return pd.DataFrame(
            columns=["qseqid", "sseqid", "pident", "length", "evalue", "bitscore"]
        )


# ========== FLASK APPLICATION ==========

app = Flask(
    __name__,
    template_folder=str(TEMPLATE_DIR),
    static_folder=str(STATIC_DIR)
)
app.secret_key = os.getenv("SECRET_KEY", "default-secret-key-change-in-production")
app.config['ENV'] = os.getenv("FLASK_ENV", "production")
app.config['DEBUG'] = os.getenv("FLASK_DEBUG", "False") == "True"

latest_predictions = []


@app.after_request
def add_cache_control(response):
    """Prevent caching of responses"""
    response.headers["Cache-Control"] = "no-cache, no-store, must-revalidate"
    response.headers["Pragma"] = "no-cache"
    response.headers["Expires"] = "0"
    return response


@app.route("/", methods=["GET", "POST"])
def index():
    """Main index page - handle sequence input and analysis"""
    global latest_predictions
    latest_predictions = []
    
    if request.method == "POST":
        sequences = []
        headers = []
        
        # Parse uploaded FASTA file
        uploaded_file = request.files.get("fasta_file")
        if uploaded_file and uploaded_file.filename != "":
            seq = ""
            header = ""
            for line in uploaded_file:
                line = line.decode().strip()
                if not line:
                    continue
                if line.startswith(">"):
                    if seq:
                        sequences.append(seq)
                        headers.append(header)
                        seq = ""
                    header = line[1:]
                else:
                    seq += line
            if seq:
                sequences.append(seq)
                headers.append(header)
        
        # Parse textarea input
        sequence_text = request.form.get("sequence_text", "").strip()
        if sequence_text:
            seq = ""
            header = ""
            for line in sequence_text.splitlines():
                line = line.strip()
                if not line:
                    continue
                if line.startswith(">"):
                    if seq:
                        sequences.append(seq)
                        headers.append(header)
                        seq = ""
                    header = line[1:]
                else:
                    seq += line
            if seq:
                sequences.append(seq)
                headers.append(header or "sequence_from_text")
        
        if not sequences:
            return render_template(
                "index.html",
                error="⚠️ No valid sequences provided.",
                current_year=datetime.now().year
            )
        
        # Detect and translate sequences
        translated_sequences = [detect_and_translate(s) for s in sequences]
        filtered_sequences = [s for s in translated_sequences if s is not None]
        filtered_headers = [
            h for s, h in zip(translated_sequences, headers) if s is not None
        ]
        
        if not filtered_sequences:
            return render_template(
                "index.html",
                error="⚠️ All sequences too short or invalid.",
                current_year=datetime.now().year
            )
        
        # Gene Family Prediction using ESM2 embeddings
        gene_preds = []
        try:
            logger.info(f"Generating embeddings for {len(filtered_sequences)} sequences...")
            features = np.array([esm2_320_embed(seq) for seq in filtered_sequences])
            
            if model and encoder:
                probs = model.predict_proba(features)
                max_probs = probs.max(axis=1)
                pred_indices = probs.argmax(axis=1)
                
                for idx, prob in zip(pred_indices, max_probs):
                    if prob >= CONFIDENCE_THRESHOLD:
                        gene_preds.append(encoder.inverse_transform([idx])[0])
                    else:
                        gene_preds.append("Unknown")
            else:
                logger.warning("Model or encoder not available, marking all as Unknown")
                gene_preds = ["Unknown"] * len(filtered_sequences)
        
        except Exception as e:
            logger.error(f"Gene family prediction failed: {e}")
            gene_preds = ["Unknown"] * len(filtered_sequences)
        
        # Save temporary FASTA for BLAST
        temp_fasta = TEMP_DIR / "temp_input.fasta"
        with open(temp_fasta, "w") as f:
            for h, s in zip(filtered_headers, filtered_sequences):
                f.write(f">{h.replace(' ', '_')}\n{s}\n")
        
        # Run BLAST
        blast_output = TEMP_DIR / "blast_results.txt"
        blast_df = run_blast_and_get_dataframe(temp_fasta, blast_output)
        
        # Load pathway mapping file
        map_df = pd.DataFrame()
        if MAPPING_FILE.exists():
            try:
                map_df = pd.read_csv(MAPPING_FILE, dtype=str)
                logger.info(f"Loaded {len(map_df)} pathway mappings")
            except Exception as e:
                logger.error(f"Failed to read mapping file: {e}")
        
        # Build results
        combined_results = []
        
        if not blast_df.empty:
            top_hits_idx = blast_df.groupby("qseqid")["bitscore"].idxmax()
            top_hits = blast_df.loc[top_hits_idx]
        else:
            top_hits = pd.DataFrame(
                columns=["qseqid", "sseqid", "pident", "length", "evalue", "bitscore"]
            )
        
        for i, orig_header in enumerate(filtered_headers):
            safe_h = orig_header.replace(" ", "_")
            top_hit_row = top_hits[top_hits["qseqid"] == safe_h]
            
            if not top_hit_row.empty:
                row = top_hit_row.iloc[0]
                top_hit = row["sseqid"]
                bitscore = float(row["bitscore"])
                pathways = "No Pathways Found"
                
                if (
                    not map_df.empty
                    and "Entry" in map_df.columns
                    and "Pathways" in map_df.columns
                ):
                    match = map_df.loc[map_df["Entry"] == top_hit, "Pathways"]
                    if not match.empty and bitscore >= BITSCORE_THRESHOLD:
                        pathways = match.values[0]
            else:
                top_hit = "No Hit"
                bitscore = 0.0
                pathways = "No Pathways Found"
            
            gene_family = gene_preds[i] if i < len(gene_preds) else "Unknown"
            combined_results.append([orig_header, top_hit, bitscore, pathways, gene_family])
        
        # Save results to CSV
        result_df = pd.DataFrame(
            combined_results,
            columns=["Query/Header", "Top Hit", "Bitscore", "Pathways", "Predicted Gene Family"]
        )
        result_df.to_csv(RESULT_CSV, index=False)
        latest_predictions = combined_results
        
        # Store in session
        session['predictions'] = combined_results
        session['result_summary'] = {
            'total_sequences': len(filtered_sequences),
            'successful_predictions': len([p for p in gene_preds if p != "Unknown"]),
            'total_pathways_found': len([r for r in combined_results if r[3] != "No Pathways Found"])
        }
        
        return redirect(url_for('results'))
    
    return render_template("index.html", predictions=[], current_year=datetime.now().year)


@app.route("/results")
def results():
    """Display analysis results"""
    predictions = session.get("predictions", [])
    summary = session.get("result_summary", {})
    
    if not predictions:
        return redirect(url_for("index"))
    
    return render_template(
        "results.html",
        predictions=predictions,
        summary=summary,
        current_year=datetime.now().year
    )


@app.route("/download_csv")
def download_csv():
    """Download results as CSV"""
    global latest_predictions
    
    if not latest_predictions:
        return "⚠️ No predictions to download."
    
    output = io.StringIO()
    writer = csv.writer(output)
    writer.writerow(["Query/Header", "Top Hit", "Bitscore", "Pathways", "Predicted Gene Family"])
    writer.writerows(latest_predictions)
    output.seek(0)
    
    return send_file(
        io.BytesIO(output.getvalue().encode()),
        mimetype="text/csv",
        as_attachment=True,
        download_name="bioinformatics_results.csv"
    )


@app.route("/health")
def health():
    """Health check endpoint"""
    return {
        "status": "healthy",
        "esm_model": esm_model is not None,
        "lr_model": model is not None,
        "encoder": encoder is not None
    }


# ========== ERROR HANDLERS ==========

@app.errorhandler(404)
def not_found(error):
    return render_template("error.html", error="Page not found"), 404


@app.errorhandler(500)
def internal_error(error):
    logger.error(f"Internal error: {error}")
    return render_template("error.html", error="Internal server error"), 500


# ========== MAIN ==========

if __name__ == "__main__":
    # Get port from environment or default to 7860 (HF Spaces)
    port = int(os.getenv("PORT", 7860))
    host = os.getenv("HOST", "0.0.0.0")
    debug = os.getenv("FLASK_DEBUG", "False") == "True"
    
    logger.info(f"Starting Flask app on {host}:{port} (debug={debug})")
    
    app.run(
        host=host,
        port=port,
        debug=debug,
        threaded=True,
        use_reloader=False
    )