import os
import re
import torch
from pathlib import Path
from contextlib import asynccontextmanager
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

# --- Configuration ---
BASE_MODEL = os.getenv("BASE_MODEL", r"Qwen/Qwen2.5-0.5B")
LORA_PATH = os.getenv("LORA_PATH", "Ibrahim-Salim1/finance-qwen-lora")
PROJECT_ROOT = Path(__file__).resolve().parents[2]
HAS_CUDA = torch.cuda.is_available()
DTYPE = torch.bfloat16 if HAS_CUDA and torch.cuda.is_bf16_supported() else torch.float16

# Global state
ml_models = {}

# --- Helper Functions (Adapted from run_chat.py) ---
def find_adapter_dir(path: Path):
    adapter_config = "adapter_config.json"
    if (path / adapter_config).is_file():
        return path
    checkpoint_dirs = sorted(
        path.glob("checkpoint-*"),
        key=lambda p: int(p.name.rsplit("-", 1)[-1]) if p.name.rsplit("-", 1)[-1].isdigit() else -1,
        reverse=True,
    )
    for checkpoint_dir in checkpoint_dirs:
        if (checkpoint_dir / adapter_config).is_file():
            return checkpoint_dir
    nested_adapters = sorted(path.rglob(adapter_config), key=lambda p: p.stat().st_mtime, reverse=True)
    if nested_adapters:
        return nested_adapters[0].parent
    return None

def resolve_adapter_path(path_str: str):
    # Support Hugging Face Hub Model IDs
    if "/" in path_str and not Path(path_str).exists() and not (PROJECT_ROOT / path_str).exists():
        return path_str
        
    path = Path(path_str).expanduser()
    candidates = []
    if path.is_absolute():
        candidates.append(path)
    else:
        candidates.extend([Path.cwd() / path, PROJECT_ROOT / path])
    
    seen = set()
    for candidate in candidates:
        candidate = candidate.resolve()
        if candidate in seen or not candidate.exists():
            continue
        seen.add(candidate)
        adapter_dir = find_adapter_dir(candidate)
        if adapter_dir is not None:
            return adapter_dir

    # Fallback search
    for search_root in [PROJECT_ROOT, Path.cwd()]:
        if not search_root.exists():
            continue
        matches = sorted(search_root.rglob("adapter_config.json"), key=lambda p: p.stat().st_mtime, reverse=True)
        if matches:
            return matches[0].parent
    return None

def build_prompt(user_prompt: str):
    return (
        "\nInstruction: You are a financial analyst. Analyze and explain clearly.\n"
        f"Input: {user_prompt}\n"
        "Response: "
    )

def clean_answer(answer: str):
    stop_texts = [
        "\nInstruction:", "\nInput:", "\nResponse:", "\nUser:", "\nAssistant:", 
        "\nQuestion:", "\nAnswer:", "chèse", "TCHA", "You are an AI assistant"
    ]
    for stop_text in stop_texts:
        if stop_text in answer:
            answer = answer.split(stop_text, 1)[0]
            
    answer = answer.strip().strip('"')
    
    # Prevent mid-sentence cutoffs when hitting max_tokens
    if answer and answer[-1] not in ".!?":
        # Find the last valid sentence boundary
        last_punct = max(answer.rfind('.'), answer.rfind('!'), answer.rfind('?'))
        if last_punct != -1:
            answer = answer[:last_punct + 1]
        else:
            answer += "..."
            
    # Remove ALL numbered and lettered prefixes
    answer = re.sub(r'(\d{1,2}|[a-z])\. ', '', answer).strip()
    
    # Replace all newlines with single spaces to form a continuous paragraph
    answer = re.sub(r'\n+', ' ', answer).strip()
    
    # Remove any isolated trailing numbers or single characters followed by a period
    answer = re.sub(r'\s+\d+\.$', '', answer).strip()
    answer = re.sub(r'\s+[a-z]\.$', '', answer).strip()
            
    return answer

# --- FastAPI Lifecycle & Endpoints ---
@asynccontextmanager
async def lifespan(app: FastAPI):
    print("Loading base model...")
    base_model = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL,
        torch_dtype=DTYPE if HAS_CUDA else torch.float32,
        device_map="auto" if HAS_CUDA else "cpu",
    )
    base_model.config.use_cache = True

    print("Loading LoRA adapter...")
    adapter_path = resolve_adapter_path(LORA_PATH)
    
    if adapter_path:
        # Always load tokenizer from the base model on Hub to ensure we get the full vocab (tokenizer.json)
        tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
        tokenizer.pad_token = tokenizer.eos_token
        model = PeftModel.from_pretrained(base_model, adapter_path)
        print(f"Loaded LoRA adapter from: {adapter_path}")
    else:
        print("Warning: Could not resolve LoRA adapter path. Using base model only.")
        tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
        tokenizer.pad_token = tokenizer.eos_token
        model = base_model
        adapter_path = "Base Model Only"

    if not HAS_CUDA:
        model = model.float() # Force all weights to float32 for CPU compatibility

    model.eval()
    device = next(model.parameters()).device

    ml_models["model"] = model
    ml_models["tokenizer"] = tokenizer
    ml_models["device"] = str(device)
    ml_models["adapter_path"] = str(adapter_path)
    
    yield
    # Clean up on shutdown
    ml_models.clear()

app = FastAPI(title="Finance Analyst API", lifespan=lifespan)

class ChatRequest(BaseModel):
    prompt: str = Field(..., description="The user's query")
    max_tokens: int = Field(60, description="Maximum number of tokens to generate")
    do_sample: bool = Field(False, description="Whether to use sampling or greedy decoding")

class ChatResponse(BaseModel):
    response: str
    
class StatusResponse(BaseModel):
    device: str
    adapter_path: str
    base_model: str

@app.get("/")
async def root():
    return {
        "message": "Finance Analyst API is running.",
        "ui_url": "http://localhost:8501",
        "docs_url": "http://localhost:8000/docs"
    }

@app.get("/status", response_model=StatusResponse)
async def get_status():
    return StatusResponse(
        device=ml_models.get("device", "unknown"),
        adapter_path=ml_models.get("adapter_path", "unknown"),
        base_model=BASE_MODEL
    )

@app.post("/chat", response_model=ChatResponse)
async def chat(request: ChatRequest):
    model = ml_models.get("model")
    tokenizer = ml_models.get("tokenizer")
    device = ml_models.get("device", "cpu")
    
    if not model or not tokenizer:
        raise HTTPException(status_code=503, detail="Model is not loaded yet")

    # Simple rule-based direct responses
    text = request.prompt.lower().strip(" .!?")
    greetings = {"hi", "hello", "hey", "salam"}
    if text in greetings:
        return ChatResponse(response="Hello. Ask me a finance question and I will help.")
        
    model_prompt = build_prompt(request.prompt)
    inputs = tokenizer(
        model_prompt,
        return_tensors="pt",
        truncation=True,
        max_length=768,
    ).to(device)

    # Handle eos_token_id potentially being a list (e.g. Qwen2.5)
    eos_id = tokenizer.eos_token_id
    pad_id = eos_id[0] if isinstance(eos_id, list) else eos_id

    generation_args = {
        "max_new_tokens": request.max_tokens,
        "do_sample": request.do_sample,
        "pad_token_id": pad_id,
        "eos_token_id": eos_id,
        "repetition_penalty": 1.05,
        "num_beams": 1,
    }
    if request.do_sample:
        generation_args.update({"temperature": 0.25, "top_p": 0.9})

    try:
        with torch.inference_mode():
            output = model.generate(**inputs, **generation_args)
            
        answer_tokens = output[0][inputs["input_ids"].shape[-1]:]
        answer = clean_answer(tokenizer.decode(answer_tokens, skip_special_tokens=True))
        
        if not answer:
            answer = "I am not sure how to answer that. Please ask a clear finance question."
    except Exception as e:
        answer = f"⚠️ Server Error during generation: {str(e)}"

    return ChatResponse(response=answer)