from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse, StreamingResponse
import os
import tempfile
from dotenv import load_dotenv
from pydantic import BaseModel
from typing import Optional, List
import anthropic
import json

# Load environment variables
load_dotenv()

# Load prompts from files
def load_prompt(filename):
    """Load prompt from text file in prompts directory"""
    try:
        prompts_dir = os.path.join(os.path.dirname(__file__), "prompts")
        file_path = os.path.join(prompts_dir, filename)
        with open(file_path, 'r', encoding='utf-8') as f:
            return f.read().strip()
    except FileNotFoundError:
        print(f"Warning: Prompt file {filename} not found. Using fallback.")
        return ""
    
def load_document(filename):
    """Load document from text file in documents directory"""
    try:
        documents_dir = os.path.join(os.path.dirname(__file__), "documents")
        file_path = os.path.join(documents_dir, filename)
        with open(file_path, 'r', encoding='utf-8') as f:
            return f.read().strip()
    except FileNotFoundError:
        print(f"Warning: Document file {filename} not found. Using fallback.")
        return ""

# Load prompts at startup
SYSTEM_PROMPT_TEMPLATE = load_prompt("system_prompt.txt")
TRANSITION_PROMPT_TEMPLATE = load_prompt("transition_prompt.txt")
DOCUMENT = load_document("cian.txt")
USER_GOAL = "More specifically, I want to understand the architecture of a transformer, why it works, and why it was designed that way. Pages 3, 4, 5, 6 are quite mysterious to me."
app = FastAPI()

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


@app.get("/api/test")
def test():
    return {"status": "OK", "app": "SokratesAI"}

class ChatMessage(BaseModel):
    role: str
    content: str
    id: Optional[str] = None

class ChatRequest(BaseModel):
    messages: List[ChatMessage]
    chunk: Optional[str] = None  # Legacy support
    currentChunk: Optional[str] = None
    nextChunk: Optional[str] = None
    action: Optional[str] = None  # 'skip', 'understood', or None
    document: Optional[str] = None
    user_goal: Optional[str] = None # User's goal for the chat, if applicable

@app.post("/api/chat")
async def chat_endpoint(request: ChatRequest):
    print(f"💬 Received chat with {len(request.messages)} messages, action: {request.action}")
    
    # Use new format if available, otherwise fall back to legacy
    current_chunk = request.currentChunk or request.chunk or "No specific chunk provided"
    next_chunk = request.nextChunk or ""
    action = request.action
    user_goal = request.user_goal or USER_GOAL

    # Only include full document on first message or transitions to provide initial context
    include_document = len(request.messages) <= 1 or action in ['skip', 'understood']
    document = DOCUMENT if include_document else ""
    
    # Create system prompt for research paper tutor with transition support
    is_transition = action in ['skip', 'understood']
    
    if is_transition:
        system_prompt = TRANSITION_PROMPT_TEMPLATE.format(
            action=action,
            current_chunk=current_chunk,
            next_chunk=next_chunk,
            document=document,
        )
    else:
        system_prompt = SYSTEM_PROMPT_TEMPLATE.format(
            current_chunk=current_chunk,
            document=document,
            user_goal=user_goal or "No specific goal provided"
        )

    anthropic_api_key = os.environ.get("ANTHROPIC_API_KEY")
    if not anthropic_api_key:
        return {"role": "assistant", "content": "I'm sorry, but the chat service is not configured. Please check the API key configuration."}

    try:
        client = anthropic.Anthropic(api_key=anthropic_api_key)

        if not request.messages:
            # No conversation yet — assistant should speak first
            anthropic_messages = [
                {"role": "user", "content": "Please start the conversation based on the provided context."}
            ]
        else:
            anthropic_messages = [
                {"role": msg.role, "content": msg.content}
                for msg in request.messages
                if msg.role in ["user", "assistant"]
            ]
            # For transitions, add a dummy user message to trigger Claude response
            if not any(msg["role"] == "user" for msg in anthropic_messages):
                if is_transition:
                    anthropic_messages.append({"role": "user", "content": "Please continue to the next section."})
                else:
                    return {"role": "assistant", "content": "I didn't receive your message. Could you please ask again?"}

        print("🤖 Calling Claude for chat response...")
        response = client.messages.create(
            model="claude-3-5-haiku-latest",
            max_tokens=10000,
            system=system_prompt,  # system prompt here
            messages=anthropic_messages,
        )

        response_text = response.content[0].text
        print(f"✅ Received response from Claude: {response_text[:100]}...")
        return {"role": "assistant", "content": response_text}

    except Exception as e:
        print(f"❌ Error in chat endpoint: {e}")
        return {"role": "assistant", "content": f"I'm sorry, I encountered an error: {str(e)}. Please try again."}


@app.post("/upload_pdf")
async def upload_pdf(file: UploadFile = File(...)):
    """Simple PDF upload endpoint that saves the file locally"""
    print(f"📄 Uploading file: {file.filename}")
    
    try:
        # Read PDF bytes
        file_bytes = await file.read()
        print(f"📊 File size: {len(file_bytes)} bytes")
        
        # Create temporary file to save PDF
        with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as temp_file:
            temp_file.write(file_bytes)
            temp_file_path = temp_file.name
        
        print(f"✅ PDF saved to: {temp_file_path}")
        
        return {
            "message": "PDF uploaded successfully!",
            "file_path": temp_file_path,
            "filename": file.filename,
            "status": "uploaded",
            "size": len(file_bytes)
        }
        
    except Exception as e:
        print(f"❌ Error uploading PDF: {e}")
        raise HTTPException(status_code=500, detail=f"PDF upload error: {str(e)}")

@app.post("/api/chat/stream")
async def chat_stream(request: ChatRequest):
    """Streaming chat endpoint for continuous conversation"""
    print(f"💬 Received chat with {len(request.messages)} messages, action: {request.action}")
    
    # Use new format if available, otherwise fall back to legacy
    current_chunk = request.currentChunk or request.chunk or "No specific chunk provided"
    next_chunk = request.nextChunk or ""
    action = request.action
    
    user_goal = request.user_goal or USER_GOAL
    
    # Only include full document on first message or transitions to provide initial context
    # After that, the conversation history maintains context
    document = DOCUMENT

    # Create system prompt for research paper tutor with transition support
    is_transition = action in ['skip', 'understood']

    print("🤖 Creating system prompt...")
    print(f"current_chunk: {current_chunk[:100] if current_chunk else 'None'}")
    print(f"next_chunk: {next_chunk[:100] if next_chunk else 'None'}")
    print(f"user_goal: {user_goal if user_goal else 'None'}")
    if is_transition:
        system_prompt = TRANSITION_PROMPT_TEMPLATE.format(
            action=action,
            current_chunk=current_chunk,
            next_chunk=next_chunk,
        )
        print(f"Transition system prompt: {system_prompt[:200]}...")
    else:
        system_prompt = SYSTEM_PROMPT_TEMPLATE.format(
            current_chunk=current_chunk,
            document=document,
            user_goal=user_goal or "No specific goal provided"
        )
        print(f"System prompt: {system_prompt[:200]}...")

    anthropic_api_key = os.environ.get("ANTHROPIC_API_KEY")
    if not anthropic_api_key:
        return {"role": "assistant", "content": "I'm sorry, but the chat service is not configured. Please check the API key configuration."}


    client = anthropic.Anthropic(api_key=anthropic_api_key)

    if not request.messages:
        # No conversation yet — assistant should speak first
        anthropic_messages = [
            {"role": "user", "content": "Please start the conversation based on the provided context."}
        ]
    else:
        anthropic_messages = [
            {"role": msg.role, "content": msg.content}
            for msg in request.messages
            if msg.role in ["user", "assistant"]
        ]
        print(anthropic_messages)
        
        # Ensure the conversation ends with a user message for thinking mode
        if anthropic_messages and anthropic_messages[-1]["role"] == "assistant":
            if is_transition:
                anthropic_messages.append({"role": "user", "content": "Please continue to the next section."})
            else:
                anthropic_messages.append({"role": "user", "content": "Please continue."})
        elif not any(msg["role"] == "user" for msg in anthropic_messages):
            if is_transition:
                anthropic_messages.append({"role": "user", "content": "Please continue to the next section."})
            else:
                def generate_error():
                    yield f"data: {json.dumps({'error': 'I did not receive your message. Could you please ask again?'})}\n\n"
                return StreamingResponse(
                    media_type="text/event-stream", 
                    content=generate_error(),
                    headers={"Cache-Control": "no-cache", 
                             "Connection": "keep-alive", 
                             "Access-Control-Allow-Origin": "*"},
                )

    def generate():
        try:
            with client.messages.stream(
                model="claude-sonnet-4-20250514",
                max_tokens=10000,
                system=system_prompt,  # system prompt here
                thinking={
                    "type": "enabled",
                    "budget_tokens": 1024,
                },
                messages=anthropic_messages,
            ) as stream:
                for text in stream.text_stream:
                    print(f"Raw text chunk: {repr(text)}")
                    yield f"data: {json.dumps(text)}\n\n"
                yield f"data: {json.dumps({'done': True})}\n\n"
        except Exception as e:
            yield f"data: {json.dumps({'error': str(e)})}\n\n"

    return StreamingResponse(
        media_type="text/event_stream", 
        content=generate(),
        headers={"Cache-Control": "no-cache", 
                 "Connection": "keep-alive", 
                 "Access-Control-Allow-Origin": "*"},
    )


# Mount static files for production deployment
frontend_path = os.path.join(os.path.dirname(__file__), "..", "frontend")
assets_path = os.path.join(frontend_path, "assets")

if os.path.exists(frontend_path):
    # Only mount assets if the directory exists (production build)
    if os.path.exists(assets_path):
        app.mount("/assets", StaticFiles(directory=assets_path), name="assets")
    
    # Serve other static files from frontend root (like pdf.worker.min.js)
    @app.get("/pdf.worker.min.js")
    async def serve_pdf_worker():
        pdf_worker_path = os.path.join(frontend_path, "pdf.worker.min.js")
        if os.path.exists(pdf_worker_path):
            return FileResponse(pdf_worker_path)
        raise HTTPException(status_code=404, detail="PDF worker not found")
    
    @app.get("/")
    async def serve_frontend():
        index_path = os.path.join(frontend_path, "index.html")
        if os.path.exists(index_path):
            return FileResponse(index_path)
        return {"message": "Backend is running - frontend not found"}
else:
    @app.get("/")
    def hello():
        return {"message": "Backend is running!"}