import os
import logging
from flask import Flask, jsonify
from flask_cors import CORS
import whisper
from dotenv import load_dotenv
from .agents.text_extractor import TextExtractorAgent
from .agents.phi_scrubber import PHIScrubberAgent
from .agents.phi_scrubber import MedicalTextUtils
from .agents.summarizer import SummarizerAgent
from .agents.medical_data_extractor import MedicalDataExtractorAgent
from .agents.medical_data_extractor import MedicalDocDataExtractorAgent
from .agents.patient_summary_agent import PatientSummarizerAgent
from .utils.model_manager import model_manager
import torch
torch.set_num_threads(1)  # Prevent PyTorch thread fighting with llama.cpp

# Load environment variables
load_dotenv()

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler('/tmp/app.log')
    ]
)

app = Flask(__name__)
CORS(app)

# Configure upload directory with safe fallbacks (avoid creating /data at import time)
def _resolve_upload_dir() -> str:
    try:
        # Prefer /data/uploads if it already exists and is writable
        data_dir = '/data/uploads'
        if os.path.isdir('/data') and (os.path.isdir(data_dir) or os.access('/data', os.W_OK)):
            os.makedirs(data_dir, exist_ok=True)
            return data_dir
    except Exception:
        pass
    # Fallback to /tmp/uploads which is always writable on Spaces
    tmp_dir = '/tmp/uploads'
    os.makedirs(tmp_dir, exist_ok=True)
    return tmp_dir

app.config['UPLOAD_FOLDER'] = _resolve_upload_dir()
app.config['MAX_CONTENT_LENGTH'] = 100 * 1024 * 1024  # 100 MB max file size

# Set cache directories
CACHE_DIRS = {
    'HF_HOME': '/tmp/huggingface',
    'XDG_CACHE_HOME': '/tmp',
    'TORCH_HOME': '/tmp/torch',
    'WHISPER_CACHE': '/tmp/whisper'
}

for env_var, path in CACHE_DIRS.items():
    os.environ[env_var] = path
    os.makedirs(path, exist_ok=True)

# WhisperModelLoader for audio transcription
class WhisperModelLoader:
    _instance = None
    
    def __init__(self):
        self._model = None
        
    @staticmethod
    def get_instance():
        if WhisperModelLoader._instance is None:
            WhisperModelLoader._instance = WhisperModelLoader()
        return WhisperModelLoader._instance
    
    def load(self):
        if self._model is None:
            try:
                logging.info("Loading Whisper model...")
                self._model = whisper.load_model(
                    "tiny",  # Using tiny model for better memory usage
                    download_root=os.environ.get('WHISPER_CACHE', '/tmp/whisper')
                )
                logging.info("Whisper model loaded successfully")
            except Exception as e:
                logging.error(f"Failed to load Whisper model: {str(e)}", exc_info=True)
                raise
        return self._model
        
    def transcribe(self, audio_path):
        model = self.load()
        return model.transcribe(audio_path)

# Initialize agents with unified model manager
try:
    # Initialize basic agents that don't require specific models
    text_extractor_agent = TextExtractorAgent()
    phi_scrubber_agent = PHIScrubberAgent()
    
    # Initialize model-dependent agents with unified model manager
    # These will be loaded dynamically when needed
    medical_data_extractor_agent = MedicalDataExtractorAgent(None)  # Will be set dynamically
    summarizer_agent = SummarizerAgent(None)  # Will be set dynamically
    
    # Initialize patient summarizer with unified model manager support
    patient_summarizer_agent = PatientSummarizerAgent(
        model_name="falconsai/medical_summarization",
        model_type="summarization"
    )
    
    # Pass all agents and models to routes
    agents = {
        "text_extractor": text_extractor_agent,
        "phi_scrubber": phi_scrubber_agent,
        "summarizer": summarizer_agent,
        "medical_data_extractor": medical_data_extractor_agent,
        "whisper_model": WhisperModelLoader.get_instance(),
        "patient_summarizer": patient_summarizer_agent,
        "model_manager": model_manager,  # Add unified model manager
    }
    
    from .api.routes import register_routes
    register_routes(app, agents)
    
    logging.info("Application initialized successfully with unified model manager")
    
except Exception as e:
    logging.error(f"Failed to initialize application: {str(e)}", exc_info=True)
    raise

@app.errorhandler(Exception)
def handle_error(error):
    logging.error(f"Unhandled error: {str(error)}", exc_info=True)
    return jsonify({
        "error": str(error),
        "status": "error"
    }), 500

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860, debug=False)