import json
import os
from core import utils
from ai import models as ai_models


def _empty_state():
    return {"active_threads": [], "immediate_handoff": "", "resolved_threads": [], "chapter": 0}


def load_story_state(folder, project_id=None):
    """Load structured story state from DB (if project_id given) or story_state.json fallback."""
    if project_id is not None:
        try:
            from web.db import StoryState
            record = StoryState.query.filter_by(project_id=project_id).first()
            if record and record.state_json:
                return json.loads(record.state_json) or _empty_state()
        except Exception:
            pass  # Fall through to file-based load if DB unavailable (e.g. CLI context)

    path = os.path.join(folder, "story_state.json")
    if os.path.exists(path):
        return utils.load_json(path) or _empty_state()
    return _empty_state()


def update_story_state(chapter_text, chapter_num, current_state, folder, project_id=None):
    """Use model_logic to extract structured story threads from the new chapter
    and save the updated state to the StoryState DB table and/or story_state.json.
    Returns the new state."""
    utils.log("STATE", f"Updating story state after Ch {chapter_num}...")
    prompt = f"""
    ROLE: Story State Tracker
    TASK: Update the structured story state based on the new chapter.

    CURRENT_STATE:
    {json.dumps(current_state)}

    NEW_CHAPTER (Chapter {chapter_num}):
    {utils.truncate_to_tokens(chapter_text, 4000)}

    INSTRUCTIONS:
    1. ACTIVE_THREADS: 2-5 concise strings, each describing what a key character is currently trying to achieve.
       - Carry forward unresolved threads from CURRENT_STATE.
       - Add new threads introduced in this chapter.
       - Remove threads that are now resolved.
    2. IMMEDIATE_HANDOFF: Write exactly 3 sentences describing how this chapter ended:
       - Sentence 1: Where are the key characters physically right now?
       - Sentence 2: What emotional state are they in at the very end of this chapter?
       - Sentence 3: What immediate unresolved threat, question, or decision is hanging in the air?
    3. RESOLVED_THREADS: Carry forward from CURRENT_STATE + add threads explicitly resolved in this chapter.

    OUTPUT_FORMAT (JSON):
    {{
        "active_threads": ["Thread 1", "Thread 2"],
        "immediate_handoff": "Sentence 1. Sentence 2. Sentence 3.",
        "resolved_threads": ["Resolved thread 1"],
        "chapter": {chapter_num}
    }}
    """
    try:
        response = ai_models.model_logic.generate_content(prompt)
        utils.log_usage(folder, ai_models.model_logic.name, response.usage_metadata)
        new_state = json.loads(utils.clean_json(response.text))
        new_state['chapter'] = chapter_num

        # Write to DB if project_id is available
        if project_id is not None:
            try:
                from web.db import db, StoryState
                from datetime import datetime
                record = StoryState.query.filter_by(project_id=project_id).first()
                if record:
                    record.state_json = json.dumps(new_state)
                    record.updated_at = datetime.utcnow()
                else:
                    record = StoryState(project_id=project_id, state_json=json.dumps(new_state))
                    db.session.add(record)
                db.session.commit()
            except Exception as db_err:
                utils.log("STATE", f"  -> DB write failed: {db_err}. Falling back to file.")

        # Always write to file for backward compat with CLI
        path = os.path.join(folder, "story_state.json")
        with open(path, 'w') as f:
            json.dump(new_state, f, indent=2)

        utils.log("STATE", f"  -> Story state saved. Active threads: {len(new_state.get('active_threads', []))}")
        return new_state
    except Exception as e:
        utils.log("STATE", f"  -> Story state update failed: {e}. Keeping previous state.")
        return current_state


def format_for_prompt(state, chapter_beats=None):
    """Format the story state into a prompt-ready string.
    Active threads and immediate handoff are always included.
    Resolved threads are only included if referenced in the chapter's beats."""
    if not state or (not state.get('immediate_handoff') and not state.get('active_threads')):
        return None

    beats_text = " ".join(str(b) for b in (chapter_beats or [])).lower()
    lines = []

    if state.get('immediate_handoff'):
        lines.append(f"IMMEDIATE STORY HANDOFF (exactly how the previous chapter ended):\n{state['immediate_handoff']}")

    if state.get('active_threads'):
        lines.append("ACTIVE PLOT THREADS:")
        for t in state['active_threads']:
            lines.append(f"  - {t}")

    relevant_resolved = [
        t for t in state.get('resolved_threads', [])
        if any(w in beats_text for w in t.lower().split() if len(w) > 4)
    ]
    if relevant_resolved:
        lines.append("RESOLVED THREADS (context only — do not re-introduce):")
        for t in relevant_resolved:
            lines.append(f"  - {t}")

    return "\n".join(lines)