import io
import torch
import torch.nn as nn
import timm
import traceback
import os
from PIL import Image
from fastapi import FastAPI, File, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from torchvision import transforms
from transformers import T5ForConditionalGeneration, T5Tokenizer
from huggingface_hub import hf_hub_download

# ─────────────────────────────────────────────────────────────────────────────
# CONFIGURATION - EXACTLY matching Colab CONFIG from SECTION 4
# ─────────────────────────────────────────────────────────────────────────────
print("="*80)
print("INITIALIZING CONFIGURATION")
print("="*80)

# Device setup - EXACTLY as Colab SECTION 3
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"PyTorch version: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"GPU Device: {torch.cuda.get_device_name(0)}")
    torch.cuda.empty_cache()
print(f"🖥️  Using device: {device}")

# Configuration - EXACTLY matching Colab SECTION 4
CONFIG = {
    # Model architecture settings
    'coatnet_model': 'coatnet_1_rw_224',
    't5_model': 't5-small',
    'img_emb_dim': 768,
    'train_last_stages': 2,

    # Image preprocessing
    'image_size': 224,

    # Inference settings
    'max_length': 100,
    'num_beams': 4,

    # Device
    'device': device
}

print("\nConfiguration loaded:")
for key, value in CONFIG.items():
    if key != 'device':
        print(f"  {key}: {value}")

# ─────────────────────────────────────────────────────────────────────────────
# SECTION 6: Model Architecture Definitions - EXACT COPY from Colab
# ─────────────────────────────────────────────────────────────────────────────
print("\n" + "="*80)
print("DEFINING MODEL ARCHITECTURES")
print("="*80)

# --- Encoder: CoAtNet --- EXACT COPY from Colab SECTION 6
class CoAtNetEncoder(nn.Module):
    def __init__(self, model_name="coatnet_1_rw_224", pretrained=True, train_last_stages=2):
        super().__init__()
        self.encoder = timm.create_model(
            model_name,
            pretrained=pretrained,
            num_classes=0,
            global_pool="avg"
        )

        # Freeze all parameters
        for p in self.encoder.parameters():
            p.requires_grad = False

        # Unfreeze last stages
        if hasattr(self.encoder, "stages") and train_last_stages is not None:
            stages = self.encoder.stages
            for stage in stages[-train_last_stages:]:
                for p in stage.parameters():
                    p.requires_grad = True

    def forward(self, x):
        return self.encoder(x)


# --- Vision-T5 Model --- EXACT COPY from Colab SECTION 6
class VisionT5Model(nn.Module):
    def __init__(self, img_encoder, txt_model_name="t5-small", img_emb_dim=768):
        super().__init__()

        # Vision encoder (CoAtNet)
        self.img_encoder = img_encoder

        # Text decoder (T5)
        self.t5 = T5ForConditionalGeneration.from_pretrained(txt_model_name)

        # Projection layer to match image features with T5 d_model
        self.proj = nn.Linear(img_emb_dim, self.t5.config.d_model)

        # Freeze shared T5 embeddings for faster and stable training
        for p in self.t5.shared.parameters():
            p.requires_grad = False

    def forward(self, pixel_values, input_ids, attention_mask, labels=None):
        # Extract image features
        img_feats = self.img_encoder(pixel_values)

        # Project image features to T5 embedding space
        img_feats = self.proj(img_feats)

        # Add sequence dimension
        encoder_hidden_states = img_feats.unsqueeze(1)

        # Run T5 encoder using image embeddings
        encoder_outputs = self.t5.encoder(
            inputs_embeds=encoder_hidden_states
        )

        # Run T5 decoder and compute loss
        outputs = self.t5(
            encoder_outputs=encoder_outputs,
            attention_mask=torch.ones(
                encoder_hidden_states.size()[:2], device=device
            ),
            input_ids=input_ids,
            labels=labels,
        )
        return outputs

    def generate_reports(self, pixel_values, max_length=100, num_beams=4):
        # Extract and project image features
        img_feats = self.img_encoder(pixel_values)
        img_feats = self.proj(img_feats)
        encoder_hidden_states = img_feats.unsqueeze(1)

        # Encode image features
        encoder_outputs = self.t5.encoder(
            inputs_embeds=encoder_hidden_states
        )

        # Generate report using beam search
        generated_ids = self.t5.generate(
            encoder_outputs=encoder_outputs,
            attention_mask=torch.ones(
                encoder_hidden_states.size()[:2], device=device
            ),
            max_length=max_length,
            num_beams=num_beams,
            early_stopping=True
        )

        return generated_ids


print("✓ Model architecture classes defined")

# ─────────────────────────────────────────────────────────────────────────────
# SECTION 7: Load Tokenizer and Image Transform - EXACT COPY from Colab
# ─────────────────────────────────────────────────────────────────────────────
print("\n" + "="*80)
print("LOADING TOKENIZER AND IMAGE TRANSFORM")
print("="*80)

# Load tokenizer
tokenizer = T5Tokenizer.from_pretrained(CONFIG['t5_model'])
print(f"✓ Loaded tokenizer: {CONFIG['t5_model']}")

# Define image transform - EXACTLY as Colab SECTION 7
transform = transforms.Compose([
    transforms.Resize((CONFIG['image_size'], CONFIG['image_size'])),
    transforms.ToTensor(),
    transforms.Normalize(
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225]
    )
])
print(f"✓ Image transform defined (size: {CONFIG['image_size']}x{CONFIG['image_size']})")

# ─────────────────────────────────────────────────────────────────────────────
# SECTION 8: Model Loading Functions - EXACT COPY from Colab
# ─────────────────────────────────────────────────────────────────────────────
def load_model_from_checkpoint(checkpoint_path: str, model_name: str, config: dict):
    """
    Load VisionT5Model from checkpoint.
    EXACT COPY from Colab SECTION 8

    Args:
        checkpoint_path: Path to .pt checkpoint file
        model_name: Name for logging (e.g., 'SFT' or 'PPO')
        config: Configuration dictionary

    Returns:
        Loaded model
    """
    print(f"\nLoading {model_name} model...")
    print(f"  Checkpoint: {checkpoint_path}")

    try:
        # Create image encoder
        print(f"  Creating CoAtNet encoder: {config['coatnet_model']}")
        img_encoder = CoAtNetEncoder(
            model_name=config['coatnet_model'],
            pretrained=False,  # Weights will come from checkpoint
            train_last_stages=config['train_last_stages']
        )

        # Create full model
        print(f"  Creating VisionT5 model with T5: {config['t5_model']}")
        model = VisionT5Model(
            img_encoder=img_encoder,
            txt_model_name=config['t5_model'],
            img_emb_dim=config['img_emb_dim']
        )

        # Load checkpoint
        print(f"  Loading checkpoint weights...")
        checkpoint = torch.load(checkpoint_path, map_location=device)

        # Handle different checkpoint formats
        if isinstance(checkpoint, dict):
            if 'model_state_dict' in checkpoint:
                state_dict = checkpoint['model_state_dict']
                print(f"  Found 'model_state_dict' in checkpoint")
            elif 'state_dict' in checkpoint:
                state_dict = checkpoint['state_dict']
                print(f"  Found 'state_dict' in checkpoint")
            elif 'model' in checkpoint:
                state_dict = checkpoint['model']
                print(f"  Found 'model' in checkpoint")
            else:
                # Assume checkpoint is the state dict
                state_dict = checkpoint
                print(f"  Using checkpoint as state_dict directly")

            # Print additional checkpoint info if available
            if 'epoch' in checkpoint:
                print(f"  Checkpoint epoch: {checkpoint['epoch']}")
            if 'loss' in checkpoint:
                print(f"  Checkpoint loss: {checkpoint['loss']:.4f}")
        else:
            state_dict = checkpoint
            print(f"  Checkpoint is a state_dict")

        # Load state dict
        missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)

        if missing_keys:
            print(f"  ⚠️ Missing keys: {len(missing_keys)}")
            if len(missing_keys) <= 5:
                for key in missing_keys:
                    print(f"    - {key}")

        if unexpected_keys:
            print(f"  ⚠️ Unexpected keys: {len(unexpected_keys)}")
            if len(unexpected_keys) <= 5:
                for key in unexpected_keys:
                    print(f"    - {key}")

        # Move to device and set to eval mode
        model = model.to(device)
        model.eval()

        print(f"✓ {model_name} model loaded successfully!")
        return model

    except Exception as e:
        print(f"❌ Error loading {model_name} model: {str(e)}")
        import traceback
        traceback.print_exc()
        raise


# ─────────────────────────────────────────────────────────────────────────────
# SECTION 9: Inference Functions - EXACT COPY from Colab
# ─────────────────────────────────────────────────────────────────────────────
def preprocess_image(image_path: str) -> torch.Tensor:
    """Load and preprocess image. EXACT COPY from Colab SECTION 9"""
    image = Image.open(image_path).convert('RGB')
    return transform(image)


def generate_report(
    image_path: str,
    model: VisionT5Model,
    config: dict
) -> str:
    """
    Generate medical report from X-ray image.
    EXACT COPY from Colab SECTION 9

    Args:
        image_path: Path to X-ray image
        model: VisionT5Model
        config: Configuration dictionary

    Returns:
        Generated report text
    """
    try:
        # Preprocess image
        pixel_values = preprocess_image(image_path).unsqueeze(0).to(device)

        # Generate report
        with torch.no_grad():
            generated_ids = model.generate_reports(
                pixel_values,
                max_length=config['max_length'],
                num_beams=config['num_beams']
            )

        # Decode
        report = tokenizer.decode(generated_ids[0], skip_special_tokens=True)

        return report.strip()

    except Exception as e:
        print(f"Error generating report for {image_path}: {str(e)}")
        return ""


# ─────────────────────────────────────────────────────────────────────────────
# LOAD MODELS FROM HUGGINGFACE
# ─────────────────────────────────────────────────────────────────────────────
print("\n" + "="*80)
print("LOADING MODELS FROM HUGGINGFACE")
print("="*80)

# Download model files from Hugging Face
try:
    SFT_MODEL_PATH = hf_hub_download(
        repo_id="vinaykumarhs2020/RLHF_radiology_model",
        filename="best_model.pt"
    )
    PPO_MODEL_PATH = hf_hub_download(
        repo_id="vinaykumarhs2020/RLHF_radiology_model",
        filename="rlhf_model.pt"
    )
    print(f"✓ Downloaded SFT model: {SFT_MODEL_PATH}")
    print(f"✓ Downloaded PPO model: {PPO_MODEL_PATH}")
except Exception as e:
    print(f"❌ Error downloading models: {e}")
    # Fallback to local paths if downloads fail
    SFT_MODEL_PATH = "/content/best_model.pt"
    PPO_MODEL_PATH = "/content/rlhf_model.pt"
    print(f"⚠️ Using local paths instead")

# Load both models - EXACTLY as Colab SECTION 8
print("\n" + "="*80)
print("LOADING MODELS")
print("="*80)

sft_model = load_model_from_checkpoint(
    SFT_MODEL_PATH,
    "SFT",
    CONFIG
)

ppo_model = load_model_from_checkpoint(
    PPO_MODEL_PATH,
    "PPO",
    CONFIG
)

print("\n✓ Both models loaded successfully!")

# ─────────────────────────────────────────────────────────────────────────────
# FASTAPI APP
# ─────────────────────────────────────────────────────────────────────────────
app = FastAPI(title="Medical Report Generation - Exact Colab Match")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.get("/health")
def health():
    return {
        "status": "ok",
        "device": str(device),
        "cuda_available": torch.cuda.is_available(),
        "models_loaded": True,
        "config": {k: v for k, v in CONFIG.items() if k != 'device'}
    }


@app.post("/sft")
async def sft_inference(file: UploadFile = File(...)):
    """
    SFT model inference - Uses EXACT generate_report() function from Colab SECTION 9
    """
    try:
        # Save uploaded file temporarily
        temp_path = f"/tmp/{file.filename}"
        with open(temp_path, "wb") as f:
            f.write(await file.read())
        
        # Use EXACT generate_report function from Colab
        report = generate_report(temp_path, sft_model, CONFIG)
        
        # Clean up temp file
        os.remove(temp_path)
        
        print(f"[SFT] Generated report: {report}")
        
        return {
            "report": report,
            "model": "SFT",
            "method": "generate_report() - exact Colab SECTION 9"
        }
        
    except Exception as e:
        traceback.print_exc()
        return {"report": f"ERROR: {str(e)}", "model": "SFT"}


@app.post("/ppo")
async def ppo_inference(file: UploadFile = File(...)):
    """
    PPO model inference - Uses EXACT generate_report() function from Colab SECTION 9
    """
    try:
        # Save uploaded file temporarily
        temp_path = f"/tmp/{file.filename}"
        with open(temp_path, "wb") as f:
            f.write(await file.read())
        
        # Use EXACT generate_report function from Colab
        report = generate_report(temp_path, ppo_model, CONFIG)
        
        # Clean up temp file
        os.remove(temp_path)
        
        print(f"[PPO] Generated report: {report}")
        
        return {
            "report": report,
            "model": "PPO",
            "method": "generate_report() - exact Colab SECTION 9"
        }
        
    except Exception as e:
        traceback.print_exc()
        return {"report": f"ERROR: {str(e)}", "model": "PPO"}


@app.post("/compare")
async def compare_models(file: UploadFile = File(...)):
    """
    Generate reports from both models for comparison
    Uses EXACT generate_report() function from Colab
    """
    try:
        # Save uploaded file temporarily
        temp_path = f"/tmp/{file.filename}"
        with open(temp_path, "wb") as f:
            f.write(await file.read())
        
        # Use EXACT generate_report function from Colab for both models
        sft_report = generate_report(temp_path, sft_model, CONFIG)
        ppo_report = generate_report(temp_path, ppo_model, CONFIG)
        
        # Clean up temp file
        os.remove(temp_path)
        
        print(f"[COMPARE] SFT: {sft_report}")
        print(f"[COMPARE] PPO: {ppo_report}")
        
        return {
            "sft_report": sft_report,
            "ppo_report": ppo_report,
            "method": "generate_report() - exact Colab SECTION 9",
            "config": {k: v for k, v in CONFIG.items() if k != 'device'}
        }
        
    except Exception as e:
        traceback.print_exc()
        return {
            "sft_report": f"ERROR: {str(e)}",
            "ppo_report": f"ERROR: {str(e)}"
        }


@app.get("/debug_inference")
def debug_inference():
    """
    Debug endpoint to verify inference setup matches Colab exactly
    """
    return {
        "device": str(device),
        "cuda_available": torch.cuda.is_available(),
        "config": {
            "coatnet_model": CONFIG['coatnet_model'],
            "t5_model": CONFIG['t5_model'],
            "img_emb_dim": CONFIG['img_emb_dim'],
            "train_last_stages": CONFIG['train_last_stages'],
            "image_size": CONFIG['image_size'],
            "max_length": CONFIG['max_length'],
            "num_beams": CONFIG['num_beams'],
        },
        "tokenizer": CONFIG['t5_model'],
        "transform": {
            "resize": f"{CONFIG['image_size']}x{CONFIG['image_size']}",
            "normalize_mean": [0.485, 0.456, 0.406],
            "normalize_std": [0.229, 0.224, 0.225]
        },
        "generation_params": {
            "max_length": CONFIG['max_length'],
            "num_beams": CONFIG['num_beams'],
            "early_stopping": True,
            "no_extra_penalties": "✓ Exactly as Colab"
        },
        "inference_method": "generate_report() from Colab SECTION 9",
        "models_loaded": {
            "sft": sft_model is not None,
            "ppo": ppo_model is not None
        },
        "model_state": {
            "sft_eval_mode": not sft_model.training if sft_model else None,
            "ppo_eval_mode": not ppo_model.training if ppo_model else None
        }
    }


# ─────────────────────────────────────────────────────────────────────────────
# STATIC FILE SERVING
# ─────────────────────────────────────────────────────────────────────────────
from fastapi.staticfiles import StaticFiles

if os.path.exists("build"):
    app.mount("/", StaticFiles(directory="build", html=True), name="static")
    print("✅ React app mounted at /")
else:
    print("⚠️ Build directory not found, serving API only")

print("\n" + "="*80)
print("SERVER READY - Using EXACT Colab Inference Code")
print("="*80)
print("Key points:")
print("  ✓ Model architecture: VisionT5Model (exact copy from Colab SECTION 6)")
print("  ✓ Inference method: generate_report() (exact copy from Colab SECTION 9)")
print("  ✓ Generation params: max_length=100, num_beams=4, early_stopping=True")
print("  ✓ No extra penalties: NO repetition_penalty, NO no_repeat_ngram_size")
print("  ✓ Transform: Resize 224x224, Normalize [0.485,0.456,0.406]/[0.229,0.224,0.225]")
print("  ✓ Device handling: Same as Colab")
print("="*80)

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