import json
import datetime
import uvicorn
from pathlib import Path
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from typing import Optional
from .model import Model
from .config import Settings

app = FastAPI(title="Shade AI Server")

# Enable CORS for external frontends (React, Flutter, etc.)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # In production, specify actual domains
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Global state for the model
active_model: Optional[Model] = None
active_settings: Optional[Settings] = None

LOGS_DIR = Path("logs")
LOGS_DIR.mkdir(exist_ok=True)

class ChatHistoryManager:
    @staticmethod
    def save_chat(model_name: str, messages: list):
        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
        log_file = LOGS_DIR / f"chat_{timestamp}.json"
        with open(log_file, "w", encoding="utf-8") as f:
            json.dump({
                "model": model_name,
                "timestamp": timestamp,
                "messages": messages
            }, f, ensure_ascii=False, indent=2)
        return log_file

class ChatRequest(BaseModel):
    message: str

class ChatResponse(BaseModel):
    response: str

# Resolve paths to the bundled web directory inside the package
WEB_DIR = Path(__file__).parent / "web"

@app.get("/api/info")
async def get_info():
    if not active_settings:
        return {"model": "No model loaded"}
    return {"model": active_settings.model}

@app.post("/api/chat")
async def chat(request: ChatRequest):
    if active_model is None or active_model.model is None:
        raise HTTPException(status_code=503, detail="Model not loaded or initialized")
    
    try:
        chat_history = [{"role": "user", "content": request.message}]
        # Using stream_chat_response which generates the string, returning as JSON matching the frontend expectation
        response_text = active_model.stream_chat_response(chat_history)
        return {"response": response_text}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/compare")
async def compare(request: ChatRequest):
    if not active_model:
        raise HTTPException(status_code=503, detail="Model not loaded")
    
    try:
        chat_history = [{"role": "user", "content": request.message}]
        
        # 1. Modified
        modified_response = active_model.stream_chat_response(chat_history)
        
        # 2. Save current LoRA weights
        lora_weights = {}
        for name, param in active_model.model.named_parameters():
            if "lora_B" in name:
                lora_weights[name] = param.data.clone()
        
        original_response = ""
        try:
            # Clear weights temporarily to act as "Original Model"
            active_model.reset_model()
            original_response = active_model.stream_chat_response(chat_history)
        finally:
            # 3. Restore weights securely
            for name, param in active_model.model.named_parameters():
                if name in lora_weights:
                    param.data.copy_(lora_weights[name])
        
        return {
            "original": original_response,
            "modified": modified_response
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/api/library")
async def get_library():
    models_dir = Path("models")
    if not models_dir.exists():
        return {"models": []}
    
    saved_models = []
    for d in models_dir.iterdir():
        if d.is_dir():
            saved_models.append({
                "name": d.name,
                "path": str(d.absolute()),
                "created": datetime.datetime.fromtimestamp(d.stat().st_ctime).isoformat()
            })
    return {"models": saved_models}

@app.get("/api/logs")
async def get_logs():
    logs = []
    if LOGS_DIR.exists():
        for f in LOGS_DIR.glob("*.json"):
            with open(f, "r", encoding="utf-8") as file:
                logs.append(json.load(file))
    return {"history": logs}

# Mount static files (JS, CSS, Images, etc.) from the build folder
if WEB_DIR.exists():
    app.mount("/", StaticFiles(directory=str(WEB_DIR), html=True), name="static")
else:
    # Log a warning if the web directory is missing
    print(f"[yellow]Warning: Web interface directory not found at {WEB_DIR}[/]")

def start_server(model_instance: Model, settings_instance: Settings, host="127.0.0.1", port=8000):
    global active_model, active_settings
    active_model = model_instance
    active_settings = settings_instance
    
    url = f"http://{host}:{port}"
    print()
    print(f"[bold green]🚀 Shade Web Server started![/]")
    print(f"[bold cyan]Click to open:[/] [blue underline]{url}[/] ")
    print()
    
    import webbrowser
    webbrowser.open(url)
    
    uvicorn.run(app, host=host, port=port, log_level="error")
