models.py

from openai import OpenAI
import json
import time
from datetime import datetime
from utils import log_step, SYSTEM_PROMPT, logger

class LSP:
    def __init__(self, api_key: str):
        if not api_key or api_key.strip() == "":
            raise ValueError("API key cannot be empty")
        self.client = OpenAI(api_key=api_key)

    def get_section_map(self, content: str) -> tuple[dict, float]:
        """Get map of sections without loading full content"""
        start_time = time.time()
        log_step("SECTION_MAP", "Starting section map generation")
        
        sections = {}
        in_code_block = False
        current_section = None
        section_start = 0
        
        lines = content.split('\n')
        for i, line in enumerate(lines):
            if line.strip().startswith('```'):
                in_code_block = not in_code_block
                continue
                
            if not in_code_block and line.strip().startswith('#'):
                if current_section:
                    sections[current_section] = (section_start, i-1)
                current_section = line.strip()
                section_start = i
                
        if current_section:
            sections[current_section] = (section_start, len(lines)-1)
        
        elapsed = time.time() - start_time
        log_step("SECTION_MAP", f"Section map generated in {elapsed:.2f}s", sections)
        return sections, elapsed

    def extract_section(self, content: str, start: int, end: int) -> tuple[str, float]:
        """Get just the content of target section"""
        start_time = time.time()
        log_step("EXTRACT", f"Extracting section from line {start} to {end}")
        
        lines = content.split('\n')
        section = '\n'.join(lines[start:end+1])
        
        elapsed = time.time() - start_time
        log_step("EXTRACT", f"Section extracted in {elapsed:.2f}s", {"length": len(section)})
        return section, elapsed

    def replace_section(self, content: str, start: int, end: int, new_content: str) -> tuple[str, float]:
        """Replace section and return full content"""
        start_time = time.time()
        log_step("REPLACE", f"Replacing section from line {start} to {end}")
        
        lines = content.split('\n')
        lines[start:end+1] = new_content.split('\n')
        result = '\n'.join(lines)
        
        elapsed = time.time() - start_time
        log_step("REPLACE", f"Section replaced in {elapsed:.2f}s", {"new_length": len(result)})
        return result, elapsed

    def tools(self):
        return [
            {
                "type": "function",
                "function": {
                    "name": "identify_section",
                    "description": "Identify which section needs to be modified based on headers map",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "section_name": {
                                "type": "string",
                                "description": "Name of section to edit"
                            },
                            "reason": {
                                "type": "string",
                                "description": "Why this section was chosen"
                            }
                        },
                        "required": ["section_name", "reason"]
                    }
                }
            },
            {
                "type": "function",
                "function": {
                    "name": "modify_section",
                    "description": "Make changes to the section content",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "modified_content": {
                                "type": "string",
                                "description": "New content for the section"
                            },
                            "changes_made": {
                                "type": "string",
                                "description": "Description of changes"
                            }
                        },
                        "required": ["modified_content", "changes_made"]
                    }
                }
            }
        ]

    async def edit_smart(self, content: str, instruction: str) -> tuple[str, list[str], float]:
        """Smart editing with section targeting"""
        total_start = time.time()
        timings = {}
        traces = []
        
        try:
            # Phase 1: Get section map
            section_map, map_time = self.get_section_map(content)
            timings['section_map'] = map_time
            traces.append(f"[{map_time:.2f}s] Found sections: {json.dumps(section_map, indent=2)}")

            # Phase 2: Identify section
            messages = [
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": f"Section map: {json.dumps(section_map)}\nInstruction: {instruction}"}
            ]

            identify_start = time.time()
            response = self.client.chat.completions.create(
                model="gpt-4o-mini",
                messages=messages,
                tools=self.tools()
            )
            identify_time = time.time() - identify_start
            timings['identify_section'] = identify_time
            
            tool_call = response.choices[0].message.tool_calls[0]
            args = json.loads(tool_call.function.arguments)
            target_section = args["section_name"]
            start, end = section_map[target_section]
            
            traces.append(f"[{identify_time:.2f}s] Selected section: {target_section} (lines {start}-{end})")
            
            # Phase 3: Extract section
            section_content, extract_time = self.extract_section(content, start, end)
            timings['extract_section'] = extract_time
            traces.append(f"[{extract_time:.2f}s] Extracted content:\n{section_content}")
            
            # Phase 4: Modify section
            modify_start = time.time()
            messages.extend([
                response.choices[0].message,
                {
                    "role": "tool",
                    "content": json.dumps({"success": True, "section": target_section}),
                    "tool_call_id": tool_call.id
                },
                {
                    "role": "user",
                    "content": f"Here's the section to modify:\n{section_content}"
                }
            ])

            response = self.client.chat.completions.create(
                model="gpt-4o-mini",
                messages=messages,
                tools=self.tools()
            )
            modify_time = time.time() - modify_start
            timings['modify_section'] = modify_time

            tool_call = response.choices[0].message.tool_calls[0]
            args = json.loads(tool_call.function.arguments)
            traces.append(f"[{modify_time:.2f}s] Modified content:\n{args['modified_content']}")

            # Phase 5: Replace section
            result, replace_time = self.replace_section(content, start, end, args["modified_content"])
            timings['replace_section'] = replace_time
            
            total_time = time.time() - total_start
            timings['total'] = total_time
            
            # Add timing summary
            timing_summary = "\nTiming Summary:\n" + "\n".join([
                f"- {step}: {time:.2f}s" for step, time in timings.items()
            ])
            traces.append(timing_summary)
            
            return result, traces, total_time

        except Exception as e:
            total_time = time.time() - total_start
            error_msg = str(e)
            traces.append(f"Error after {total_time:.2f}s: {error_msg}")
            return error_msg, traces, total_time

    async def edit_naive(self, content: str, instruction: str) -> tuple[str, list[str], float]:
        """Naive approach - give everything to AI"""
        total_start = time.time()
        traces = [f"[0.00s] Starting naive edit (sending entire document)"]
        
        try:
            messages = [
                {"role": "system", "content": "You are a document editor. Edit the provided document according to instructions."},
                {"role": "user", "content": f"Edit this document according to this instruction: {instruction}\n\nDocument:\n{content}"}
            ]

            api_start = time.time()
            response = self.client.chat.completions.create(
                model="gpt-4o-mini",
                messages=messages
            )
            api_time = time.time() - api_start

            result = response.choices[0].message.content
            total_time = time.time() - total_start
            
            traces.extend([
                f"[{api_time:.2f}s] OpenAI API call completed",
                f"[{total_time:.2f}s] Total processing completed"
            ])
            
            return result, traces, total_time
            
        except Exception as e:
            total_time = time.time() - total_start
            error_msg = str(e)
            traces.append(f"Error after {total_time:.2f}s: {error_msg}")
            return error_msg, traces, total_time