import os
import xml.etree.ElementTree as ET
from flask import Flask, request, send_from_directory
from litellm import completion
from functools import lru_cache

app = Flask(__name__, static_folder=os.path.join(os.path.dirname(__file__), 'static'))

# Get version from environment or use default
SERVICE_VERSION = os.environ.get('COPILOT_VERSION', "copilot-service v1.0")
DEFAULT_MODEL = os.environ.get('DEFAULT_MODEL', "gpt-4.1")
DEFAULT_MAX_TOKENS = int(os.environ.get('DEFAULT_MAX_TOKENS', "50"))

app.config["DEFAULT_SYSTEM_PROMPT"] = """
            You are an expert coding prediction model. You are to receive the contents of a file, along with the location of the user's cursor. The location shall be marked with █
            Then you are to predict the next line of code that the user is likely to write.
            You are to return the next line of code in XML format, like this:
            <completion>rint("Hello World!")</completion>
            You may also perform multi-line completions, like this:
            <multiline-completion>llo_world():
            print("Hello World!")
            </multiline-completion>
            YOU MUST ONLY RETURN THE XML. DO NOT RETURN ANY OTHER TEXT.
            Only perform multiline completions if the user is likely to write more than one line of code, for example, if the user is writing a function or a class.

            EXAMPLES:
            1. If you see: "using System;█" (C#), return: <completion>using System.Collections.Generic;</completion>
            2. If you see: "function fetchData() █" (JavaScript), return: <completion>{ return fetch('/api'); }</completion>
            3. If you see: "class User █" (PHP), return: <multiline-completion>{\n    protected $name;\n    \n    public function __construct($name) {\n        $this->name = $name;\n    }\n}</multiline-completion>
            4. If you see: "SELECT * F█" (SQL), return: <completion>FROM users WHERE id = 1;</completion>
            5. If you see: "def calculate█" (Python), return: <multiline-completion>_sum(a, b):\n    return a + b</multiline-completion>
            6. If you see: "public static void m█" (Java), return: <completion>ain(String[] args) {</completion>
            7. If you see极狐: "interface User █" (TypeScript), return: <multiline-completion>{\n  id: number;\n  name: string;\n}</multiline-complication>
            8. If you see: "func handleErr█" (Go), return: <multiline-completion>or(w http.ResponseWriter, status int) {\n    w.WriteHeader(status)\n}</multiline-completion>
            9. If you see: "Rails.applic█" (Ruby), return: <multiline-completion>tion.routes.draw do\n  resources :users\nend</multiline-completion>
            10. If you see: "enum HttpSt█" (Rust), return: <multiline-completion>atus {\n    Ok,\n    NotFound,\n}</multiline-completion>
""".strip()

app.config["DEFAULT_USER_PROMPT"] = """
            You are an expert coding prediction model. You are to receive the contents of a file, along with the location of the user's cursor. The location shall be marked with █
            Then you are to predict the next line of code that the user is likely to write.
            You are to return the next line of code in XML format. Only return XML.
            Only perform multiline completions when appropriate.

            EXAMPLES:
            1. Input: "using System;█" (C#) → Output: <completion>using System.Collections.Generic;</completion>
            2. Input: "function fetchData() █" (JavaScript) → Output: <completion>{ return fetch('/api'); }</completion>
            3. Input: "class User █" (PHP) → Output: <multiline-completion>{\n    protected $name;\n    \n    public function __construct($name) {\n        $this->name = $name;\n    }\n}</multiline-completion>
            4. Input: "SELECT * F█" (SQL) → Output: <completion>ROM users WHERE id = 1;</completion>
            5. Input: "def calculate█" (Python) → Output: <multiline-completion>_sum(a, b):\n    return a + b</multiline-completion>
            6. Input: "public static void m█" (Java) → Output: <completion>ain(String[] args) {</completion>
            7. Input: "interface User █" (TypeScript) → Output: <multiline-completion>{\n  id: number;\n  name: string;\n}</multiline-completion>
            8. Input: "func handleErr█" (Go) → Output: <multiline-completion>or(w http.ResponseWriter, status int) {\n    w.WriteHeader(status)\n}</multiline-completion>
            9. Input: "Rails.applic█" (Ruby) → Output: <multiline-completion>tion.routes.draw do\n  resources :users\nend</multiline-completion>
            10. Input: "enum HttpSt█" (Rust) → Output: <multiline-completion>atus {\n    Ok,\n    NotFound,\n}</multiline-completion>

            |||FILE_CONTENTS_PLACEHOLDER|||
""".strip()


@lru_cache(maxsize=100)
def cached_completion(model: str, system_prompt: str, user_prompt: str, max_tokens: int):
    """Cached version of the LiteLLM completion call"""
    return completion(
        model=model,
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt}
        ],
        max_tokens=max_tokens
    )

def parseXML(xmlfile):
    # Use regex to get the content of the <completion> tha
    completion = ""
    try:
        root = ET.fromstring(xmlfile)
        print(f"Root tag: {root.tag}")
        print(f"Root text: {root.text}")
        completion = root.text

    except ET.ParseError as e:
        raise ValueError(f"Invalid XML format: {e}")
    return completion


""" Endpoint to handle the completions  requests
This endpoint expects a JSON body with the following structure:
{ 
    "file_contents": "The contents of the file as a string",
    "cursor_position": 42  # The position of the cursor in所述 file contents. This must be an integer, based off tbe index of the character in the string.
}
"""
@app.route("/copilot/completions", methods=["POST"])
def create_completion():
    body = request.get_json()
    print(body)
    if not body:
        print("No body provided in the request")
        return {"error": "Invalid request, body is required"}, 400
    if "file_contents" not in body:
        print("No file_contents provided in the request")
        return {"error": "Invalid request, 'file_contents' is required"}, 400
    if "cursor_position" not in body:
        print("No cursor_position provided in the request")
        return {"error": "Invalid request, 'cursor_position' is required"}, 400
    if not isinstance(body["file_contents"], str):
        print("file_contents is not a string")
        return {"error": "Invalid request, 'file_contents' must be a string"}, 400
    if not isinstance(body["cursor_position"], int):
        print("cursor_position is not an integer")
        return {"error": "Invalid request, 'cursor_position' must be an integer"}, 400
    if not (0 <= body["cursor_position"] <= len(body["file_contents"])):
        print("cursor_position is out of bounds")
        return {"error": "Invalid request, 'cursor_position' must be within the bounds of 'file_contents'"}, 400

    file_contents = body.get("file_contents", "")
    cursor_position = body.get("cursor_position", 0)
    content_list = list(file_contents)
    content_list.insert(cursor_position, "█")
    file_contents_with_cursor_position = ''.join(content_list)

    # Get overrides
    model_override = body.get("model", DEFAULT_MODEL)
    max_tokens_override = body.get("max_tokens_override", DEFAULT_MAX_TOKENS)
    system_prompt_override = body.get("system_prompt_override")
    user_prompt_override = body.get("user_prompt_override")

    system_content = system_prompt_override or app.config["DEFAULT_SYSTEM_PROMPT"]
    if user_prompt_override:
        user_content = user_prompt_override
    else:
        base_prompt = app.config["DEFAULT_USER_PROMPT"]
        user_content = base_prompt.replace(
            "|||FILE_CONTENTS_PLACEHOLDER|||",
            file_contents_with_cursor_position
        )

    response = cached_completion(
        model_override,
        system_content,
        user_content,
        max_tokens_override
    )
    if not response or not response.choices:
        return {"error": "No completion returned from the model"}, 500
    try:
        xml_content = response.choices[0].message.content.strip()
        print(f"XML Content: {xml_content}")
        completion_result = parseXML(xml_content)
        print(f"Parsed Completion: {completion_result}")
        
        # Add token accounting metadata
        token_usage = {
            "prompt_tokens": response.usage.get("prompt_tokens", 0),
            "completion_tokens": response.usage.get("completion_tokens", 0),
            "total_tokens": response.usage.get("total_tokens", 0)
        }
        return {
            "completion": completion_result,
            "token_metadata": token_usage,
            "model_used": model_override
        }, 200
    except ValueError as e:
        return {"error": str(e)}, 400


@app.route("/version", methods=["GET"])
def version():
    return SERVICE_VERSION, 200

@app.route("/copilot/preview", methods=["POST"])
def preview_prompt():
    body = request.get_json()
    # Reuse validation from completion endpoint
    if not body:
        return {"error": "Invalid request, body is required"}, 400
    if "file_contents" not in body:
        return {"error": "Invalid request, 'file_contents' is required"}, 400
    if "cursor_position" not in body:
        return {"error": "Invalid request, 'cursor_position' is required"}, 400
    if not isinstance(body["file_contents"], str):
        return {"error": "Invalid request, 'file_contents' must be a string"}, 400
    if not isinstance(body["cursor_position"], int):
        return {"error": "Invalid request, 'cursor_position' must be an integer"}, 400

    file_contents = body["file_contents"]
    cursor_position = body["cursor_position"]
    if not (0 <= cursor_position <= len(file_contents)):
        return {"error": "Invalid request, 'cursor_position' must be within the bounds of 'file_contents'"}, 400

    # Prepare file content with cursor
    content_list = list(file_contents)
    content_list.insert(cursor_position, "█")
    file_contents_with_cursor = ''.join(content_list)

    # Apply overrides
    system_prompt = body.get("system_prompt_override") or app.config.get("DEFAULT_SYSTEM_PROMPT")
    if user_prompt_override := body.get("user_prompt_override"):
        user_prompt = user_prompt_override
    else:
        base_prompt = app.config["DEFAULT_USER_PROMPT"]
        user_prompt = base_prompt.replace(
            "|||FILE_CONTENTS_PLACEHOLDER|||",
            file_contents_with_cursor
        )

    return {
        "system_prompt": system_prompt,
        "user_prompt": user_prompt,
        "max_tokens": body.get("max_tokens_override", DEFAULT_MAX_TOKENS)
    }, 200

@app.route('/')
def serve_docs():
    return send_from_directory('static', 'index.html')

@app.route('/static/<path:path>')
def serve_static(path):
    return send_from_directory('static', path)

def main():
    port = int(os.environ.get('PORT', 5000))
    app.run(host="0.0.0.0", port=port)

if __name__ == "__main__":
    main()
