import os
import json
import time
from typing import Any, Dict, Optional, Union
import torch
import torch.nn as nn
from transformers import PreTrainedModel, AutoModelForCausalLM, AutoTokenizer
from safetensors.torch import load_file
from peft import PeftModel
from .configuration_kutup import KutupConfig

class EvidenceRoutingBlock(nn.Module):
    def __init__(self, d_model: int, n_heads: int = 8, dropout: float = 0.05):
        super().__init__()
        self.cross_attn = nn.MultiheadAttention(d_model, n_heads, dropout=dropout, batch_first=True)
        self.norm = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, query: torch.Tensor, key_value: torch.Tensor) -> torch.Tensor:
        attn_out, _ = self.cross_attn(query, key_value, key_value)
        return self.norm(query + self.dropout(attn_out))

class JointSchemaHead(nn.Module):
    def __init__(self, hidden_size: int = 2048, decision_width: int = 1024, num_layers: int = 4):
        super().__init__()
        self.proj = nn.Linear(hidden_size, decision_width)
        self.routing = EvidenceRoutingBlock(decision_width)
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=decision_width,
            nhead=8,
            dim_feedforward=decision_width * 4,
            dropout=0.05,
            activation="gelu",
            batch_first=True,
            norm_first=True,
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        self.scorer = nn.Sequential(
            nn.Linear(decision_width, decision_width // 2),
            nn.GELU(),
            nn.Linear(decision_width // 2, 1),
        )

    def forward(self, prefix_emb: torch.Tensor, option_embs: torch.Tensor) -> torch.Tensor:
        p_proj = self.proj(prefix_emb)
        o_proj = self.proj(option_embs)
        routed = self.routing(o_proj, p_proj)
        joint = self.transformer(routed)
        scores = self.scorer(joint).squeeze(-1)
        return scores

class KutupForDecision(PreTrainedModel):
    config_class = KutupConfig

    def __init__(self, config: KutupConfig):
        super().__init__(config)
        self.config = config
        self.head = JointSchemaHead(
            hidden_size=config.hidden_size,
            decision_width=config.decision_width
        )

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
        config = kwargs.pop("config", None)
        if config is None:
            config = KutupConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)
        
        device_map = kwargs.get("device_map", None)
        torch_dtype = getattr(torch, config.torch_dtype) if isinstance(config.torch_dtype, str) else torch.bfloat16
        
        model = cls(config)
        
        # 1. Backbone Yükle
        print(f"[KUTUP-1] Backbone yukleniyor: {config.backbone_model_id}...")
        backbone = AutoModelForCausalLM.from_pretrained(
            config.backbone_model_id,
            torch_dtype=torch_dtype,
            device_map=device_map,
        )
        
        # 2. LoRA Varsa Yukle
        lora_dir = os.path.join(pretrained_model_name_or_path, "lora_adapter")
        if os.path.exists(lora_dir):
            print(f"[KUTUP-1] LoRA yukleniyor: {lora_dir}...")
            backbone = PeftModel.from_pretrained(backbone, lora_dir)
            backbone = backbone.merge_and_unload()
        model.backbone = backbone
        
        # 3. Head safetensors yukle
        head_path = os.path.join(pretrained_model_name_or_path, "joint_head.safetensors")
        if os.path.exists(head_path):
            print(f"[KUTUP-1] Joint Head yukleniyor: {head_path}...")
            weights = load_file(head_path)
            model.head.load_state_dict(weights, strict=False)
            
        model.to(torch_dtype)
        model.eval()
        return model

    @torch.inference_mode()
    def decide(self, tokenizer: AutoTokenizer, state: str, questions: Dict[str, Any]) -> Dict[str, Any]:
        device = next(self.parameters()).device
        answers = {}
        for qid, qdef in questions.items():
            qtype = qdef.get("type", "choice")
            inst = qdef.get("instructions", "")
            crit = qdef.get("criteria", {})
            
            prefix_text = f"[DURUM]
{state}

[SORU]
{inst}"
            p_inputs = tokenizer(prefix_text, return_tensors="pt", truncation=True, max_length=self.config.max_length).to(device)
            p_out = self.backbone.model(**p_inputs)
            p_emb = p_out.last_hidden_state
            
            opts = []
            if qtype == "noul":
                opts = ["Evet", "Hayır"]
                opt_keys = ["true", "false"]
            elif isinstance(crit, dict):
                opt_keys = list(crit.keys())
                opts = [f"{k}: {v}" for k, v in crit.items()]
            elif isinstance(crit, list):
                opt_keys = [str(i) for i in range(len(crit))]
                opts = [f"Seviye {i}: {v}" for i, v in enumerate(crit)]
                
            opt_embs = []
            for opt in opts:
                o_inp = tokenizer(f"[SECENEK] {opt}", return_tensors="pt").to(device)
                o_out = self.backbone.model(**o_inp)
                opt_embs.append(o_out.last_hidden_state.mean(dim=1))
                
            all_opt_emb = torch.stack(opt_embs, dim=1) # [1, num_opts, hidden]
            logits = self.head(p_emb, all_opt_emb).squeeze(0) / self.config.temperature
            probs = torch.softmax(logits, dim=-1).tolist()
            
            prob_dict = {k: round(p, 4) for k, p in zip(opt_keys, probs)}
            best_idx = int(torch.argmax(logits).item())
            best_key = opt_keys[best_idx]
            conf = round(probs[best_idx], 4)
            
            if qtype == "noul":
                answers[qid] = {"type": "noul", "noul": prob_dict.get("true", 0.5), "confidence": conf}
            else:
                answers[qid] = {"choice": best_key, "confidence": conf, "probabilities": prob_dict}
                
        return {"model": "TurkishCodeMan/KUTUP-1-Decision-2B", "answers": answers}