# app.py

from typing import Optional, Dict, List, Tuple
import gradio as gr

from constants import HTML_SYSTEM_PROMPT, AVAILABLE_MODELS, DEMO_LIST
from hf_client import get_inference_client
from tavily_search import enhance_query_with_search
from utils import (
    extract_text_from_file,
    extract_website_content,
    apply_search_replace_changes,
    history_to_messages,
    history_to_chatbot_messages,
    remove_code_block,
    parse_transformers_js_output,
    format_transformers_js_output
)
from deploy import send_to_sandbox, load_project_from_url

# Type aliases
History = List[Tuple[str, str]]

# Core generation function
def generation_code(
    query: Optional[str],
    image: Optional[gr.Image],
    file: Optional[str],
    website_url: Optional[str],
    _setting: Dict[str, str],
    _history: Optional[History],
    _current_model_name: str,
    enable_search: bool,
    language: str,
    provider: str
) -> Tuple[str, History, str, List[Dict[str, str]]]:
    # Initialize inputs
    if query is None:
        query = ''
    if _history is None:
        _history = []

    # System prompt and history
    system_prompt = _setting.get('system', HTML_SYSTEM_PROMPT)
    messages = history_to_messages(_history, system_prompt)

    # File input
    if file:
        text = extract_text_from_file(file)
        query += f"\n\n[File content]\n{text[:5000]}"

    # Website input
    if website_url:
        text = extract_website_content(website_url)
        if not text.startswith('Error'):
            query += f"\n\n[Website content]\n{text[:8000]}"

    # Web search enhancement
    final_query = enhance_query_with_search(query, enable_search)
    messages.append({'role': 'user', 'content': final_query})

    # Model inference
    client = get_inference_client(_current_model_name, provider)
    resp = client.chat.completions.create(
        model=_current_model_name,
        messages=messages,
        max_tokens=10000
    )
    content = resp.choices[0].message.content

    # Post-processing
    has_existing = bool(_history)
    if language == 'transformers.js':
        files = parse_transformers_js_output(content)
        code_str = format_transformers_js_output(files)
        preview_html = send_to_sandbox(files['index.html'])
    else:
        clean = remove_code_block(content)
        if has_existing:
            clean = apply_search_replace_changes(_history[-1][1], clean)
        code_str = clean
        preview_html = send_to_sandbox(clean) if language == 'html' else ''

    # Update history
    new_history = _history + [(query, code_str)]
    chat_msgs = history_to_chatbot_messages(new_history)

    return code_str, new_history, preview_html, chat_msgs

# Build UI
with gr.Blocks(theme=gr.themes.Base(), title="AnyCoder - AI Code Generator") as demo:
    # State
    history_state = gr.State([])
    setting_state = gr.State({'system': HTML_SYSTEM_PROMPT})
    model_state = gr.State(AVAILABLE_MODELS[0]['id'])

    with gr.Sidebar():
        gr.Markdown("## AnyCoder AI")

        # Load project
        url_in = gr.Textbox(label="Load HF Space URL", placeholder="https://huggingface.co/spaces/user/project")
        load_btn = gr.Button("Import Project")
        load_status = gr.Markdown(visible=False)

        gr.Markdown("---")

        # Inputs
        prompt_in = gr.Textbox(label="Prompt", lines=3)
        file_in = gr.File(label="Reference file")
        image_in = gr.Image(label="Design image")
        url_site = gr.Textbox(label="Website URL")
        search_chk = gr.Checkbox(label="Enable Web Search")
        language_dd = gr.Dropdown(
            choices=["html", "python", "transformers.js"],
            value="html",
            label="Language"
        )
        model_dd = gr.Dropdown(
            choices=[m['name'] for m in AVAILABLE_MODELS],
            value=AVAILABLE_MODELS[0]['name'],
            label="Model"
        )
        gen_btn = gr.Button("Generate")
        clr_btn = gr.Button("Clear")

    with gr.Column():
        with gr.Tabs():
            with gr.Tab("Code"):
                code_out = gr.Code(label="Generated Code")
            with gr.Tab("Preview"):
                preview_out = gr.HTML(label="Live Preview")
            with gr.Tab("History"):
                chat_out = gr.Chatbot(label="History")

    # Events
    load_btn.click(
        fn=lambda u: load_project_from_url(u),
        inputs=[url_in],
        outputs=[load_status, code_out, preview_out, url_in, history_state, chat_out]
    )

    def on_model_change(name: str) -> str:
        for m in AVAILABLE_MODELS:
            if m['name'] == name:
                return m['id']
        return AVAILABLE_MODELS[0]['id']

    model_dd.change(
        fn=on_model_change,
        inputs=[model_dd],
        outputs=[model_state]
    )

    gen_btn.click(
        fn=generation_code,
        inputs=[
            prompt_in, image_in, file_in, url_site,
            setting_state, history_state, model_state,
            search_chk, language_dd, gr.State('auto')
        ],
        outputs=[code_out, history_state, preview_out, chat_out]
    )

    clr_btn.click(
        fn=lambda: ([], [], "", []),
        outputs=[history_state, chat_out, preview_out, code_out]
    )

if __name__ == '__main__':
    demo.queue().launch()